1 files changed, 32 insertions, 12 deletions

diff --git a/pkg/libs/cancellablewebsocket/cancellablewebsocket.go b/pkg/libs/cancellablewebsocket/cancellablewebsocket.go
index 5dcffc1..accc180 100644
--- a/pkg/libs/cancellablewebsocket/cancellablewebsocket.go
+++ b/pkg/libs/cancellablewebsocket/cancellablewebsocket.go
@@ -9,10 +9,11 @@ import (
 )
 
 type CancellableWebSocket struct {
-	conn   *websocket.Conn
-	ctx    context.Context
-	cancel context.CancelFunc
-	errors chan error
+	conn      *websocket.Conn
+	ctx       context.Context
+	cancel    context.CancelFunc
+	closeCode int // 0 means killed
+	errors    chan error
 }
 
 func Dial(dialer *websocket.Dialer, ctx context.Context, url string, requestHeader http.Header) (*CancellableWebSocket, error) {
@@ -24,10 +25,11 @@ func Dial(dialer *websocket.Dialer, ctx context.Context, url string, requestHead
 	childCtx, cancel := context.WithCancel(ctx)
 
 	cws := &CancellableWebSocket{
-		conn:   conn,
-		ctx:    childCtx,
-		cancel: cancel,
-		errors: make(chan error),
+		conn:      conn,
+		ctx:       childCtx,
+		cancel:    cancel,
+		closeCode: 0,
+		errors:    make(chan error),
 	}
 
 	go cws.listenForCancel()
@@ -57,7 +59,7 @@ func (cws *CancellableWebSocket) WriteJSON(v interface{}) error {
 	return nil
 }
 
-func (cws *CancellableWebSocket) Close() error {
+func (cws *CancellableWebSocket) Kill() error {
 	cws.cancel()
 	err, ok := <-cws.errors
 
@@ -68,8 +70,16 @@ func (cws *CancellableWebSocket) Close() error {
 	return nil
 }
 
-func (cws *CancellableWebSocket) listenForCancel() {
+func (cws *CancellableWebSocket) Close(code int) error {
+	cws.closeCode = code
+	if err := cws.Kill(); err != nil {
+		cws.closeCode = 0 // Reset close code
+		return err
+	}
+	return nil
+}
 
+func (cws *CancellableWebSocket) listenForCancel() {
 	<-cws.ctx.Done()
 
 	var err error
@@ -83,8 +93,18 @@ func (cws *CancellableWebSocket) listenForCancel() {
 		cws.errors <- err
 	}
 
-	if err = cws.conn.Close(); err != nil {
-		cws.errors <- err
+	if cws.closeCode == 0 {
+		// Close without sending close.
+		if err = cws.conn.Close(); err != nil {
+			cws.errors <- err
+		}
+	} else {
+		// Send close with code.
+		// NOTE: The SetWriteDeadline above does not affect this. :)
+		err = cws.conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(cws.closeCode, ""), time.Now().Add(time.Second))
+		if err != nil {
+			cws.errors <- err
+		}
 	}
 
 	close(cws.errors)