auth.go

package middleware

import (
	"encoding/json"
	"io"
	"net/http"

	"github.com/gin-gonic/gin"
	"github.com/rs/zerolog/log"
	"gitlab.com/tensorsecurity-rd/waf-console/internal/utils"
)

const (
	// Cookie name to check for the auth token
	AuthCookieName = "auth_token"
)

// AuthMiddleware validates the auth cookie with SSO service
func AuthMiddleware(ssoUrl string) gin.HandlerFunc {
	return func(c *gin.Context) {
		// skip ping
		if c.Request.URL.Path != "/ping" {
			c.Next()
			return
		}
		// Get auth cookie
		cookies := c.Request.Cookies()
		if len(cookies) == 0 {
			utils.AssembleResponse(c, nil, utils.ErrUnauthorized)
			c.Abort()
			return
		}

		// Create request to SSO service
		req, err := http.NewRequest(http.MethodPost, ssoUrl, nil)
		if err != nil {
			utils.AssembleResponse(c, nil, utils.ErrInternalServer)
			c.Abort()
			return
		}

		// Add auth cookie to request
		for _, cookie := range cookies {
			req.AddCookie(cookie)
		}

		// Make request to SSO service
		client := &http.Client{}
		resp, err := client.Do(req)
		if err != nil {
			utils.AssembleResponse(c, nil, utils.ErrInternalServer)
			c.Abort()
			return
		}
		defer resp.Body.Close()

		// Read response body
		body, err := io.ReadAll(resp.Body)
		if err != nil {
			log.Error().Err(err).Msg("failed to read sso response body")
			utils.AssembleResponse(c, nil, utils.ErrInternalServer)
			c.Abort()
			return
		}

		// Parse SSO response
		var ssoResp SSOResponse
		if err := json.Unmarshal(body, &ssoResp); err != nil {
			log.Error().Err(err).Msg("failed to unmarshal sso response")
			utils.AssembleResponse(c, nil, utils.ErrInternalServer)
			c.Abort()
			return
		}

		// Check if authentication was successful
		if ssoResp.Code != "OK" {
			utils.AssembleResponse(c, nil, utils.ErrUnauthorized)
			c.Abort()
			return
		}

		// Store user info in context for later use
		c.Set("userInfo", ssoResp.Data)

		// Authentication successful, continue
		c.Next()
	}
}