Fix sending on closed channels (#1)

Author: LucaScorpion

pkg/proxy/proxy.go

@@ -90,7 +90,9 @@ func (p *Proxy) onAgentConnected(newClient *ws.Client) {
 	go func() {
 		for msg := range newAgent.client.Recv {
 			for r := range newAgent.remotes {
-				r.client.Send <- msg
+				if !r.client.Closed() {
+					r.client.Send <- msg
+				}
 			}
 		}
 
@@ -132,7 +134,9 @@ func (p *Proxy) onRemoteConnected(newClient *ws.Client) {
 	// Pipe received messages to the agent.
 	go func() {
 		for msg := range newRemote.client.Recv {
-			targetAgent.client.Send <- msg
+			if !targetAgent.client.Closed() {
+				targetAgent.client.Send <- msg
+			}
 		}
 
 		log.Println("Closed remote:", target.Id)

pkg/ws/client.go

@@ -33,6 +33,10 @@ func (c *Client) Close() {
 	}
 }
 
+func (c *Client) Closed() bool {
+	return c.closed
+}
+
 func (c *Client) readPump() {
 	defer c.Close()
 
@@ -52,6 +56,10 @@ func (c *Client) readPump() {
 			break
 		}
 
+		if c.closed {
+			break
+		}
+
 		c.Recv <- Message{
 			MessageType: t,
 			Data:        msg,

pkg/ws/server.go

@@ -13,8 +13,8 @@ type Server struct {
 }
 
 var upgrader = websocket.Upgrader{
+	// Allow all origins.
 	CheckOrigin: func(_ *http.Request) bool {
-		// TODO: Actually check origin
 		return true
 	},
 }