rpi-vsts-ci-led/pkg/websocket/websocket.go
2020-04-17 16:41:48 +08:00

92 lines
1.5 KiB
Go

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
}