Merge pull request #38 from shatteredsilicon/fix-connection-hijacked-…

…problem

Fix 'http: connection has been hijacked' problem

app/controllers/agent/ws.go

@@ -66,7 +66,7 @@ func (c Agent) Cmd(uuid string, conn *websocket.Conn) revel.Result {
 	// When the agent disconnects, set oN.agent_configs.running=0 for the agent.
 	dbm := c.Args["dbm"].(db.Manager)
 	if err := dbm.Open(); err != nil {
-		return c.Error(err, "Agent.Cmd: dbm.Open")
+		return c.RenderError(fmt.Errorf("Agent.Cmd: dbm.Open: %s", err.Error()))
 	}
 	agentHandler := agent.NewMySQLHandler(dbm, instance.NewMySQLHandler(dbm))
 
@@ -122,7 +122,7 @@ func (c Agent) Data(conn *websocket.Conn) revel.Result {
 	dbStats := msgStats // copy
 	dbm := db.NewMySQLManager()
 	if err := dbm.Open(); err != nil {
-		return c.Error(err, "Agent.Data: dbm.Open")
+		return c.RenderError(fmt.Errorf("Agent.Data: dbm.Open: %s", err.Error()))
 	}
 	defer dbm.Close()
 
@@ -141,7 +141,7 @@ func (c Agent) Data(conn *websocket.Conn) revel.Result {
 		case io.EOF:
 			// We got everything, client disconnected.
 		default:
-			return c.Error(err, "Agent.Data: qan.SaveData")
+			return c.RenderError(fmt.Errorf("Agent.Data: qan.SaveData: %s", err.Error()))
 		}
 	}
 
@@ -160,7 +160,7 @@ func (c Agent) Log(conn *websocket.Conn) revel.Result {
 	dbStats := msgStats // copy
 	dbm := db.NewMySQLManager()
 	if err := dbm.Open(); err != nil {
-		return c.Error(err, "Agent.Log: dbm.Open")
+		return c.RenderError(fmt.Errorf("Agent.Log: dbm.Open: %s", err.Error()))
 	}
 	defer dbm.Close()
 	dbh := agent.NewLogHandler(dbm, &dbStats)
@@ -176,7 +176,7 @@ func (c Agent) Log(conn *websocket.Conn) revel.Result {
 		case io.EOF:
 			revel.TRACE.Printf("%s: done (EOF)", prefix)
 		default:
-			return c.Error(err, "Agent.Log: agent.SaveLog")
+			return c.RenderError(fmt.Errorf("Agent.Log: agent.SaveLog: %s", err.Error()))
 		}
 	}
 

app/init.go

@@ -40,6 +40,7 @@ import (
 	queryService "github.com/shatteredsilicon/qan-api/service/query"
 	"github.com/shatteredsilicon/qan-api/stats"
 	"github.com/shatteredsilicon/ssm/proto"
+	"golang.org/x/net/websocket"
 
 	// dummy import to keep the package
 	_ "github.com/revel/modules/testrunner/app/controllers"
@@ -250,7 +251,16 @@ func authAgent(c *revel.Controller) revel.Result {
 			revel.ERROR.Printf("auth agent: %s", err)
 		}
 		c.Response.Status = int(res.Code)
-		return c.RenderText(res.Error)
+		if c.Request.Method == "WS" {
+			if err := websocket.JSON.Send(c.Request.Websocket, proto.Response{
+				Code:  res.Code,
+				Error: res.Error,
+			}); err != nil {
+				return c.RenderError(err)
+			}
+		} else {
+			return c.RenderText(res.Error)
+		}
 	}
 	c.Args["agentId"] = agentId