core/cookies.go

// Copyright (c) 2024 Zeni Kim <zenik@smarteching.com>
// Use of this source code is governed by MIT-style
// license that can be found in the LICENSE file.

package core

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"encoding/gob"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"strconv"
	"strings"
)

var (
	ErrValueTooLong = errors.New("cookie value too long")
	ErrInvalidValue = errors.New("invalid cookie value")
)

// Declare the User type.
type UserCookie struct {
	Email string
	Token string
}

var secretcookie []byte

func GetCookie(r *http.Request) (UserCookie, error) {

	var err error
	// Create a new instance of a User type.
	var user UserCookie

	// check if template engine is enable
	TemplateEnableStr := os.Getenv("TEMPLATE_ENABLE")
	if TemplateEnableStr == "" {
		TemplateEnableStr = "false"
	}
	TemplateEnable, _ := strconv.ParseBool(TemplateEnableStr)
	// if enabled,
	if TemplateEnable {
		cookie_secret := os.Getenv("COOKIE_SECRET")
		if cookie_secret == "" {
			panic("cookie secret key is not set")
		}
		secretcookie, err = hex.DecodeString(cookie_secret)
		if err != nil {
			return user, err
		}
	} else {
		panic("Templates are disabled")
	}

	gobEncodedValue, err := CookieReadEncrypted(r, "goffee", secretcookie)
	if err != nil {
		return user, err
	}

	// Create an strings.Reader containing the gob-encoded value.
	reader := strings.NewReader(gobEncodedValue)

	// Decode it into the User type. Notice that we need to pass a *pointer* to
	// the Decode() target here?
	if err := gob.NewDecoder(reader).Decode(&user); err != nil {
		return user, err
	}

	return user, nil

}

func SetCookie(w http.ResponseWriter, email string, token string) error {
	// Initialize a User struct containing the data that we want to store in the
	// cookie.
	var err error

	// check if template engine is enable
	TemplateEnableStr := os.Getenv("TEMPLATE_ENABLE")
	if TemplateEnableStr == "" {
		TemplateEnableStr = "false"
	}
	TemplateEnable, _ := strconv.ParseBool(TemplateEnableStr)
	// if enabled,
	if TemplateEnable {
		cookie_secret := os.Getenv("COOKIE_SECRET")
		if cookie_secret == "" {
			panic("cookie secret key is not set")
		}
		secretcookie, err = hex.DecodeString(cookie_secret)
		if err != nil {
			return err
		}
	} else {
		panic("Templates are disabled")
	}

	if err != nil {
		return err
	}

	user := UserCookie{Email: email, Token: token}

	// Initialize a buffer to hold the gob-encoded data.
	var buf bytes.Buffer

	// Gob-encode the user data, storing the encoded output in the buffer.
	err = gob.NewEncoder(&buf).Encode(&user)
	if err != nil {
		return err
	}

	// Call buf.String() to get the gob-encoded value as a string and set it as
	// the cookie value.
	cookie := http.Cookie{
		Name:     "goffee",
		Value:    buf.String(),
		Path:     "/",
		MaxAge:   3600,
		HttpOnly: true,
		SameSite: http.SameSiteLaxMode,
	}

	// Write an encrypted cookie containing the gob-encoded data as normal.
	err = CookieWriteEncrypted(w, cookie, secretcookie)
	if err != nil {
		return err
	}

	return nil
}

func CookieWrite(w http.ResponseWriter, cookie http.Cookie) error {
	// Encode the cookie value using base64.
	cookie.Value = base64.URLEncoding.EncodeToString([]byte(cookie.Value))

	// Check the total length of the cookie contents. Return the ErrValueTooLong
	// error if it's more than 4096 bytes.
	if len(cookie.String()) > 4096 {
		return ErrValueTooLong
	}

	// Write the cookie as normal.
	http.SetCookie(w, &cookie)

	return nil
}

func CookieRead(r *http.Request, name string) (string, error) {
	// Read the cookie as normal.
	cookie, err := r.Cookie(name)
	if err != nil {
		return "", err
	}

	// Decode the base64-encoded cookie value. If the cookie didn't contain a
	// valid base64-encoded value, this operation will fail and we return an
	// ErrInvalidValue error.
	value, err := base64.URLEncoding.DecodeString(cookie.Value)
	if err != nil {
		return "", ErrInvalidValue
	}

	// Return the decoded cookie value.
	return string(value), nil
}

func CookieWriteEncrypted(w http.ResponseWriter, cookie http.Cookie, secretKey []byte) error {
	// Create a new AES cipher block from the secret key.
	block, err := aes.NewCipher(secretKey)
	if err != nil {
		return err
	}

	// Wrap the cipher block in Galois Counter Mode.
	aesGCM, err := cipher.NewGCM(block)
	if err != nil {
		return err
	}

	// Create a unique nonce containing 12 random bytes.
	nonce := make([]byte, aesGCM.NonceSize())
	_, err = io.ReadFull(rand.Reader, nonce)
	if err != nil {
		return err
	}

	// Prepare the plaintext input for encryption. Because we want to
	// authenticate the cookie name as well as the value, we make this plaintext
	// in the format "{cookie name}:{cookie value}". We use the : character as a
	// separator because it is an invalid character for cookie names and
	// therefore shouldn't appear in them.
	plaintext := fmt.Sprintf("%s:%s", cookie.Name, cookie.Value)

	// Encrypt the data using aesGCM.Seal(). By passing the nonce as the first
	// parameter, the encrypted data will be appended to the nonce — meaning
	// that the returned encryptedValue variable will be in the format
	// "{nonce}{encrypted plaintext data}".
	encryptedValue := aesGCM.Seal(nonce, nonce, []byte(plaintext), nil)

	// Set the cookie value to the encryptedValue.
	cookie.Value = string(encryptedValue)

	// Write the cookie as normal.
	return CookieWrite(w, cookie)
}

func CookieReadEncrypted(r *http.Request, name string, secretKey []byte) (string, error) {
	// Read the encrypted value from the cookie as normal.
	encryptedValue, err := CookieRead(r, name)
	if err != nil {
		return "", err
	}

	// Create a new AES cipher block from the secret key.
	block, err := aes.NewCipher(secretKey)
	if err != nil {
		return "", err
	}

	// Wrap the cipher block in Galois Counter Mode.
	aesGCM, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}

	// Get the nonce size.
	nonceSize := aesGCM.NonceSize()

	// To avoid a potential 'index out of range' panic in the next step, we
	// check that the length of the encrypted value is at least the nonce
	// size.
	if len(encryptedValue) < nonceSize {
		return "", ErrInvalidValue
	}

	// Split apart the nonce from the actual encrypted data.
	nonce := encryptedValue[:nonceSize]
	ciphertext := encryptedValue[nonceSize:]

	// Use aesGCM.Open() to decrypt and authenticate the data. If this fails,
	// return a ErrInvalidValue error.
	plaintext, err := aesGCM.Open(nil, []byte(nonce), []byte(ciphertext), nil)
	if err != nil {
		return "", ErrInvalidValue
	}

	// The plaintext value is in the format "{cookie name}:{cookie value}". We
	// use strings.Cut() to split it on the first ":" character.
	expectedName, value, ok := strings.Cut(string(plaintext), ":")
	if !ok {
		return "", ErrInvalidValue
	}

	// Check that the cookie name is the expected one and hasn't been changed.
	if expectedName != name {
		return "", ErrInvalidValue
	}

	// Return the plaintext cookie value.
	return value, nil
}