grpc/auth/session.go

package auth

import (
	"context"
	"fmt"
	"net/http"

	"github.com/gorilla/securecookie"
	"github.com/gorilla/sessions"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
	"google.golang.org/protobuf/proto"
)

// TokenResponse is a proto message that contains a token.
// It typically is a response from a login service.
type TokenResponse interface {
	proto.Message
	GetToken() string
}

// Handler is an interface that can be used to handle token forwarding.
// It can be used to forward tokens from grpc messages to cookies and vice versa.
// T is the token response type that creates the session and contains the token to be forwarded as a cookie.
// U is the message that closes the session.
type Handler interface {
	CookieToAuth(ctx context.Context) context.Context
	ForwardTokenAsCookie(ctx context.Context, w http.ResponseWriter, p proto.Message) error
	ForwardCookieAsToken(ctx context.Context, request *http.Request) metadata.MD
	SendAuthCookie(ctx context.Context, token string) error
}

func NewHandler(sessionName string, sessions *sessions.CookieStore) Handler {
	return &handler{sessionName: sessionName, sessions: sessions}
}

type handler struct {
	sessionName string
	sessions    *sessions.CookieStore
}

// ForwardTokenAsCookie forwards the token from the message to the cookie store.
func (h *handler) ForwardTokenAsCookie(_ context.Context, w http.ResponseWriter, p proto.Message) error {
	switch m := any(p).(type) {
	case TokenResponse:
		sess := h.newSession(h.sessionName)
		sess.Values["token"] = m.GetToken()
		// request is not available in the context, but session.Save does not actually use it
		return h.sessions.Save(nil, w, sess)
	default:
		sess := h.newSession(h.sessionName)
		sess.Options.MaxAge = -1
		return h.sessions.Save(nil, w, sess)
	}
}

// ForwardCookieAsToken forwards the token from the cookie store to the grpc metadata.
func (h *handler) ForwardCookieAsToken(_ context.Context, request *http.Request) metadata.MD {
	sess, err := h.sessions.Get(request, h.sessionName)
	if err != nil {
		return nil
	}
	tk, ok := sess.Values["token"]
	if !ok {
		return nil
	}
	return metadata.Pairs("authorization", fmt.Sprintf("Bearer %v", tk))
}

// SendAuthCookie sets the cookie in the response.
func (h *handler) SendAuthCookie(ctx context.Context, token string) error {
	sess := h.newSession(h.sessionName)
	sess.Values["token"] = token
	if token == "" {
		sess.Options.MaxAge = -1
	}
	encoded, err := securecookie.EncodeMulti(sess.Name(), sess.Values, h.sessions.Codecs...)
	if err != nil {
		return err
	}
	return grpc.SetHeader(ctx, metadata.Pairs("set-cookie", sessions.NewCookie(sess.Name(), encoded, sess.Options).String()))
}

func (h *handler) newSession(name string) *sessions.Session {
	session := sessions.NewSession(h.sessions, name)
	opts := *h.sessions.Options
	session.Options = &opts
	session.IsNew = true
	return session
}

// CookieToAuth forwards the cookie authorization to the grpc metadata.
func (h *handler) CookieToAuth(ctx context.Context) context.Context {
	md, ok := metadata.FromIncomingContext(ctx)
	if !ok {
		return ctx
	}
	// if we have an authorization header, we don't need to check for cookies here
	if len(md.Get("authorization")) > 0 {
		return ctx
	}
	// no other way to easily parse cookies
	header := http.Header{}
	for _, v := range md.Get("cookie") {
		header.Add("Cookie", v)
	}
	md.Delete("cookie")
	sess, err := h.sessions.Get(&http.Request{Header: header}, h.sessionName)
	if err != nil {
		return ctx
	}
	if v, ok := sess.Values["token"]; ok {
		md["authorization"] = []string{fmt.Sprintf("Bearer %s", v)}
	}
	return metadata.NewIncomingContext(ctx, md)
}