rig-go-sdk/interceptor.go at main · rigdev/rig-go-sdk · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
package rig

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

	"connectrpc.com/connect"
	"github.com/golang-jwt/jwt"
	"github.com/rigdev/rig-go-api/api/v1/authentication"
	"github.com/rigdev/rig-go-api/api/v1/authentication/authenticationconnect"
)

var _omitAuth = map[string]struct{}{
	"/api.v1.authentication.Service/Login":             {},
	"/api.v1.authentication.Service/Register":          {},
	"/api.v1.authentication.Service/VerifyEmail":       {},
	"/api.v1.authentication.Service/RefreshToken":      {},
	"/api.v1.authentication.Service/OauthCallback":     {},
	"/api.v1.authentication.Service/SendPasswordReset": {},
	"/api.v1.authentication.Service/ResetPassword":     {},
	"/api.v1.authentication.Service/GetAuthConfig":     {},
}

type authInterceptor struct {
	cfg *config
}

func (i *authInterceptor) handleAuth(ctx context.Context, h http.Header, method string) {
	if _, ok := _omitAuth[method]; !ok {
		if i.cfg.basicAuth != "" {
			i.setBasicAuth(ctx, h)
		} else {
			i.setAuthorization(ctx, h)
		}
	}
}

func (i *authInterceptor) setBasicAuth(ctx context.Context, h http.Header) {
	h.Set("Authorization", i.cfg.basicAuth)
}

func (i *authInterceptor) setAuthorization(ctx context.Context, h http.Header) {
	at := i.cfg.sm.GetAccessToken()
	if at == "" && i.cfg.login != nil {
		res, err := authenticationconnect.NewServiceClient(i.cfg.hc, i.cfg.host).Login(ctx, &connect.Request[authentication.LoginRequest]{
			Msg: i.cfg.login,
		})
		if err != nil {
			return
		}

		i.cfg.sm.SetAccessToken(res.Msg.GetToken().GetAccessToken(), res.Msg.GetToken().GetRefreshToken())
	}

	c := jwt.StandardClaims{}
	p := jwt.Parser{
		SkipClaimsValidation: true,
	}

	_, _, err := p.ParseUnverified(i.cfg.sm.GetAccessToken(), &c)
	if err != nil {
		return
	}

	if !c.VerifyExpiresAt(time.Now().Add(30*time.Second).Unix(), true) {
		res, err := authenticationconnect.NewServiceClient(i.cfg.hc, i.cfg.host).RefreshToken(ctx, &connect.Request[authentication.RefreshTokenRequest]{
			Msg: &authentication.RefreshTokenRequest{
				RefreshToken: i.cfg.sm.GetRefreshToken(),
			},
		})
		if err != nil {
			return
		}
		i.cfg.sm.SetAccessToken(res.Msg.GetToken().GetAccessToken(), res.Msg.GetToken().GetRefreshToken())
	}

	h.Set("Authorization", fmt.Sprint("Bearer ", i.cfg.sm.GetAccessToken()))
}

func (i *authInterceptor) WrapUnary(next connect.UnaryFunc) connect.UnaryFunc {
	return func(ctx context.Context, ar connect.AnyRequest) (connect.AnyResponse, error) {
		i.handleAuth(ctx, ar.Header(), ar.Spec().Procedure)
		return next(ctx, ar)
	}
}

func (i *authInterceptor) WrapStreamingClient(next connect.StreamingClientFunc) connect.StreamingClientFunc {
	return func(ctx context.Context, s connect.Spec) connect.StreamingClientConn {
		conn := next(ctx, s)
		i.handleAuth(ctx, conn.RequestHeader(), s.Procedure)
		return conn
	}
}

func (i *authInterceptor) WrapStreamingHandler(next connect.StreamingHandlerFunc) connect.StreamingHandlerFunc {
	return func(ctx context.Context, shc connect.StreamingHandlerConn) error {
		i.handleAuth(ctx, shc.RequestHeader(), shc.Spec().Procedure)
		return next(ctx, shc)
	}
}