woodpecker/router/middleware/session/user.go

package session

import (
	"net/http"

	"github.com/drone/drone/model"
	"github.com/drone/drone/router/middleware/context"
	"github.com/drone/drone/shared/token"

	"github.com/gin-gonic/gin"
)

func User(c *gin.Context) *model.User {
	v, ok := c.Get("user")
	if !ok {
		return nil
	}
	u, ok := v.(*model.User)
	if !ok {
		return nil
	}
	return u
}

func Token(c *gin.Context) *token.Token {
	v, ok := c.Get("token")
	if !ok {
		return nil
	}
	u, ok := v.(*token.Token)
	if !ok {
		return nil
	}
	return u
}

func SetUser() gin.HandlerFunc {
	return func(c *gin.Context) {
		var user *model.User

		t, err := token.ParseRequest(c.Request, func(t *token.Token) (string, error) {
			var db = context.Database(c)
			var err error
			user, err = model.GetUserLogin(db, t.Text)
			return user.Hash, err
		})
		if err == nil {
			c.Set("user", user)

			// if this is a session token (ie not the API token)
			// this means the user is accessing with a web browser,
			// so we should implement CSRF protection measures.
			if t.Kind == token.SessToken {
				err = token.CheckCsrf(c.Request, func(t *token.Token) (string, error) {
					return user.Hash, nil
				})
				// if csrf token validation fails, exit immediately
				// with a not authorized error.
				if err != nil {
					c.AbortWithStatus(http.StatusUnauthorized)
					return
				}
			}
		}
		c.Next()
	}
}

func MustAdmin() gin.HandlerFunc {
	return func(c *gin.Context) {
		user := User(c)
		switch {
		case user == nil:
			c.AbortWithStatus(http.StatusUnauthorized)
			// c.HTML(http.StatusUnauthorized, "401.html", gin.H{})
		case user.Admin == false:
			c.AbortWithStatus(http.StatusForbidden)
			// c.HTML(http.StatusForbidden, "401.html", gin.H{})
		default:
			c.Next()
		}

	}
}

func MustUser() gin.HandlerFunc {
	return func(c *gin.Context) {
		user := User(c)
		switch {
		case user == nil:
			c.AbortWithStatus(http.StatusUnauthorized)
			// c.HTML(http.StatusUnauthorized, "401.html", gin.H{})
		default:
			c.Next()
		}

	}
}