// Vikunja is a to-do list application to facilitate your life.
// Copyright 2018-2021 Vikunja and contributors. All rights reserved.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public Licensee as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public Licensee for more details.
//
// You should have received a copy of the GNU Affero General Public Licensee
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

package identityawareproxy

import (
	"fmt"
	"sync"
	"time"

	"code.vikunja.io/api/pkg/config"
	"code.vikunja.io/api/pkg/log"
	"code.vikunja.io/api/pkg/modules/auth"
	"code.vikunja.io/web/handler"
	"github.com/dgrijalva/jwt-go"
	"github.com/labstack/echo/v4"
	"github.com/lestrrat-go/jwx/jwk"
)

// TimeFunc provides the current time to help validate "exp" and "iss" claims in a JWT.
// It is overridden in unit tests
var TimeFunc = time.Now

// Caches the public keys of the identity-aware proxy used to validate the auth data it sends
type iapCache struct {
	keyset *jwk.Set
	mutex  sync.Mutex
}

// GetKeyset returns the cached public keys from the identity-aware proxy
// or fetches them for the first time.
func (cache *iapCache) GetKeyset() (*jwk.Set, error) {
	if cache.keyset != nil {
		return cache.keyset, nil
	}

	cache.mutex.Lock()
	defer cache.mutex.Unlock()

	// Check that another thread has not fetched the keyset
	if cache.keyset != nil {
		return cache.keyset, nil
	}

	// Fetch the public key(s) from the identity-aware proxy
	keyset, err := jwk.FetchHTTP(config.AuthIdentityAwareProxyJwksURI.GetString())
	if err != nil {
		log.Error("Failed to retrieve the identity-aware proxy's signing public key at URL %s: %v", config.AuthIdentityAwareProxyJwksURI.GetString(), err)
		return nil, ErrIAPPublicKeysetMissing{URL: config.AuthIdentityAwareProxyJwksURI.GetString()}
	}
	cache.keyset = keyset
	return cache.keyset, nil
}

// The identity-aware proxy authentication middleware parses and validates the
// JWT provided by the IAP
func Middleware() echo.MiddlewareFunc {
	cache := &iapCache{}
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			// Skip if IAP is not enabled
			if !config.AuthIdentityAwareProxyEnabled.GetBool() {
				return next(c)
			}

			// Get and validate claims from the IAP
			token := c.Request().Header.Get(config.AuthIdentityAwareProxyJwtHeader.GetString())
			if token == "" {
				err := ErrIAPTokenMissing{Header: config.AuthIdentityAwareProxyJwtHeader.GetString()}
				return handler.HandleHTTPError(err, c)
			}

			keyset, err := cache.GetKeyset()
			if err != nil {
				return handler.HandleHTTPError(err, c)
			}

			cl, err := parseAndValidateJwt(token, keyset)
			if err != nil {
				return handler.HandleHTTPError(err, c)
			}
			c.Set(IAPClaimsContextKey, cl)

			// Generate auth.AuthClaims from the IAP identity
			user, err := userForIAPClaims(cl)
			if err != nil {
				return handler.HandleHTTPError(err, c)
			}
			authClaims := newIAPUserJWTAuthClaims(user)
			c.Set(auth.AuthClaimsContextKey, authClaims)
			return next(c)
		}
	}
}

// The authMiddleware generates and stores internal auth based on
// those claims. This overwrites any auth from the JWT middleware

func parseAndValidateJwt(token string, keyset *jwk.Set) (*IAPClaims, error) {
	// Parse the jwt from the identity-aware proxy using the correct key
	tken, err := jwt.ParseWithClaims(token, &IAPClaims{}, func(unvalidatedToken *jwt.Token) (interface{}, error) {
		// Only support either ECDSA or RSA signing methods. Never support the "none" signing method
		if _, ok := unvalidatedToken.Method.(*jwt.SigningMethodECDSA); !ok {
			if _, ok := unvalidatedToken.Method.(*jwt.SigningMethodRSA); !ok {
				return nil, ErrIAPUnsupportedJWTSigningMethod{Method: unvalidatedToken.Header["alg"].(string)}
			}
		}

		keyID, ok := unvalidatedToken.Header["kid"].(string)
		if !ok {
			return nil, ErrIAPJWTMissingKID{}
		}

		keys := keyset.LookupKeyID(keyID)
		if len(keys) != 1 {
			return nil, ErrIAPJWTMissingKID{}
		}

		var rawkey interface{} // This is the raw key, like *rsa.PublicKey or *ecdsa.PublicKey
		if err := keys[0].Raw(&rawkey); err != nil {
			return nil, err
		}

		return jwk.PublicKeyOf(rawkey)
	})
	if err != nil {
		return nil, err
	}

	cl, ok := tken.Claims.(*IAPClaims)
	if !ok || !tken.Valid {
		return nil, fmt.Errorf("failed to parse the jwt claims")
	}

	return cl, nil
}