add new

pkg/websocket/websocket.go

@@ -0,0 +1,91 @@
+package websocket
+
+import (
+	"errors"
+	"net/http"
+	"net/textproto"
+
+	"github.com/gorilla/websocket"
+)
+
+type Socket struct {
+	*websocket.Conn
+	Listeners map[string](chan<- []byte)
+	done      chan struct{}
+	errd      bool
+	errsig    chan error
+}
+
+var client *Socket
+
+func NewClient(remote string) (*Socket, error) {
+	if remote == "" {
+		return nil, errors.New("remote url is empty")
+	}
+	var err error
+
+	headers := http.Header{}
+
+	headers.Set(
+		textproto.CanonicalMIMEHeaderKey("X-Auth-Key"),
+		"rpi",
+	)
+	c, resp, err := websocket.DefaultDialer.Dial(remote, headers)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	client = &Socket{
+		Conn:      c,
+		Listeners: make(map[string](chan<- []byte)),
+		done:      make(chan struct{}, 0),
+		errd:      false,
+		errsig:    make(chan error),
+	}
+
+	return client, nil
+}
+
+func (p *Socket) Listen() error {
+	go func() {
+		defer close(p.done)
+		for {
+			_, msg, err := p.ReadMessage()
+			if err != nil {
+				p.errsig <- err
+			}
+
+			for _, v := range p.Listeners {
+				v <- msg
+			}
+		}
+	}()
+
+	select {
+	case <-p.done:
+		return nil
+	case err := <-p.errsig:
+		return err
+	}
+}
+
+func (p *Socket) AddListener(alias string, c chan<- []byte) error {
+	if _, ok := p.Listeners[alias]; ok {
+		return errors.New("listener name exists")
+	}
+	p.Listeners[alias] = c
+
+	return nil
+}
+
+func (p *Socket) RemoveListener(alias string) error {
+	if v, ok := p.Listeners[alias]; !ok {
+		return errors.New("listener not exists")
+	} else {
+		close(v)
+		delete(p.Listeners, alias)
+	}
+
+	return nil
+}