github.com/Jeffail/benthos/v3@v3.65.0/lib/output/writer/websocket.go (about) 1 package writer 2 3 import ( 4 "crypto/tls" 5 "net/http" 6 "net/url" 7 "sync" 8 "time" 9 10 "github.com/Jeffail/benthos/v3/lib/log" 11 "github.com/Jeffail/benthos/v3/lib/metrics" 12 "github.com/Jeffail/benthos/v3/lib/types" 13 "github.com/Jeffail/benthos/v3/lib/util/http/auth" 14 btls "github.com/Jeffail/benthos/v3/lib/util/tls" 15 "github.com/gorilla/websocket" 16 ) 17 18 //------------------------------------------------------------------------------ 19 20 // WebsocketConfig contains configuration fields for the Websocket output type. 21 type WebsocketConfig struct { 22 URL string `json:"url" yaml:"url"` 23 auth.Config `json:",inline" yaml:",inline"` 24 TLS btls.Config `json:"tls" yaml:"tls"` 25 } 26 27 // NewWebsocketConfig creates a new WebsocketConfig with default values. 28 func NewWebsocketConfig() WebsocketConfig { 29 return WebsocketConfig{ 30 URL: "ws://localhost:4195/post/ws", 31 Config: auth.NewConfig(), 32 TLS: btls.NewConfig(), 33 } 34 } 35 36 //------------------------------------------------------------------------------ 37 38 // Websocket is an output type that serves Websocket messages. 39 type Websocket struct { 40 log log.Modular 41 stats metrics.Type 42 43 lock *sync.Mutex 44 45 conf WebsocketConfig 46 client *websocket.Conn 47 tlsConf *tls.Config 48 } 49 50 // NewWebsocket creates a new Websocket output type. 51 func NewWebsocket( 52 conf WebsocketConfig, 53 log log.Modular, 54 stats metrics.Type, 55 ) (*Websocket, error) { 56 ws := &Websocket{ 57 log: log, 58 stats: stats, 59 lock: &sync.Mutex{}, 60 conf: conf, 61 } 62 if conf.TLS.Enabled { 63 var err error 64 if ws.tlsConf, err = conf.TLS.Get(); err != nil { 65 return nil, err 66 } 67 } 68 return ws, nil 69 } 70 71 //------------------------------------------------------------------------------ 72 73 func (w *Websocket) getWS() *websocket.Conn { 74 w.lock.Lock() 75 ws := w.client 76 w.lock.Unlock() 77 return ws 78 } 79 80 //------------------------------------------------------------------------------ 81 82 // Connect establishes a connection to an Websocket server. 83 func (w *Websocket) Connect() error { 84 w.lock.Lock() 85 defer w.lock.Unlock() 86 87 if w.client != nil { 88 return nil 89 } 90 91 headers := http.Header{} 92 93 purl, err := url.Parse(w.conf.URL) 94 if err != nil { 95 return err 96 } 97 98 if err := w.conf.Sign(&http.Request{ 99 URL: purl, 100 Header: headers, 101 }); err != nil { 102 return err 103 } 104 105 var client *websocket.Conn 106 if w.conf.TLS.Enabled { 107 dialer := websocket.Dialer{ 108 TLSClientConfig: w.tlsConf, 109 } 110 if client, _, err = dialer.Dial(w.conf.URL, headers); err != nil { 111 return err 112 113 } 114 } else if client, _, err = websocket.DefaultDialer.Dial(w.conf.URL, headers); err != nil { 115 return err 116 } 117 118 go func(c *websocket.Conn) { 119 for { 120 if _, _, cerr := c.NextReader(); cerr != nil { 121 c.Close() 122 break 123 } 124 } 125 }(client) 126 127 w.client = client 128 return nil 129 } 130 131 //------------------------------------------------------------------------------ 132 133 // Write attempts to write a message by pushing it to an Websocket broker. 134 func (w *Websocket) Write(msg types.Message) error { 135 client := w.getWS() 136 if client == nil { 137 return types.ErrNotConnected 138 } 139 140 err := msg.Iter(func(i int, p types.Part) error { 141 return client.WriteMessage(websocket.BinaryMessage, p.Get()) 142 }) 143 if err != nil { 144 w.lock.Lock() 145 w.client = nil 146 w.lock.Unlock() 147 if err == websocket.ErrCloseSent { 148 return types.ErrNotConnected 149 } 150 return err 151 } 152 return nil 153 } 154 155 // CloseAsync shuts down the Websocket output and stops processing messages. 156 func (w *Websocket) CloseAsync() { 157 go func() { 158 w.lock.Lock() 159 if w.client != nil { 160 w.client.Close() 161 w.client = nil 162 } 163 w.lock.Unlock() 164 }() 165 } 166 167 // WaitForClose blocks until the Websocket output has closed down. 168 func (w *Websocket) WaitForClose(timeout time.Duration) error { 169 return nil 170 } 171 172 //------------------------------------------------------------------------------