From 779c1c06d3fcf56864e4da60808e64dab1276552 Mon Sep 17 00:00:00 2001
From: Brad Rydzewski <brad.rydzewski@gmail.com>
Date: Thu, 30 Apr 2015 14:57:53 -0700
Subject: [PATCH] agent endpoints now protected

---
 drone.go         |  1 +
 server/agent.go  |  2 ++
 server/server.go | 15 +++++++++++++++
 3 files changed, 18 insertions(+)

diff --git a/drone.go b/drone.go
index c4f80c7f..c396b62f 100644
--- a/drone.go
+++ b/drone.go
@@ -115,6 +115,7 @@ func main() {
 
 	queue := api.Group("/queue")
 	{
+		queue.Use(server.MustAgent())
 		queue.GET("", server.GetQueue)
 		queue.POST("/pull", server.PollBuild)
 
diff --git a/server/agent.go b/server/agent.go
index 702e32e8..609dd888 100644
--- a/server/agent.go
+++ b/server/agent.go
@@ -9,6 +9,8 @@ import (
 func GetAgentToken(c *gin.Context) {
 	sess := ToSession(c)
 	token := &common.Token{}
+	token.Kind = common.TokenAgent
+	token.Label = "drone-agent"
 	tokenstr, err := sess.GenerateToken(token)
 	if err != nil {
 		c.Fail(500, err)
diff --git a/server/server.go b/server/server.go
index bb0cb0cd..1a7c692c 100644
--- a/server/server.go
+++ b/server/server.go
@@ -213,6 +213,21 @@ func MustAdmin() gin.HandlerFunc {
 	}
 }
 
+func MustAgent() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		sess := ToSession(c)
+		token := sess.GetLogin(c.Request)
+		if token == nil {
+			c.AbortWithStatus(401)
+			return
+		} else if token.Kind != common.TokenAgent {
+			c.AbortWithStatus(500)
+			return
+		}
+		c.Next()
+	}
+}
+
 func CheckPull() gin.HandlerFunc {
 	return func(c *gin.Context) {
 		u := ToUser(c)