package ext import ( "context" "encoding/base64" "errors" "net/http" "time" "github.com/sirupsen/logrus" "git.sr.ht/~gabrielgio/img/pkg/database/repository" "git.sr.ht/~gabrielgio/img/pkg/service" ) func HTML(next http.HandlerFunc) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/html") next(w, r) } } type LogMiddleware struct { entry *logrus.Entry } func NewLogMiddleare(log *logrus.Entry) *LogMiddleware { return &LogMiddleware{ entry: log, } } func (l *LogMiddleware) HTTP(next http.HandlerFunc) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { start := time.Now() next(w, r) elapsed := time.Since(start) l.entry. WithField("time", elapsed). WithField("path", r.URL.Path). Info(r.Method) } } type AuthMiddleware struct { key []byte entry *logrus.Entry } func NewAuthMiddleware(key []byte, log *logrus.Entry) *AuthMiddleware { return &AuthMiddleware{ key: key, entry: log.WithField("context", "auth"), } } func (a *AuthMiddleware) LoggedIn(next http.HandlerFunc) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { path := string(r.URL.Path) if path == "/login" || path == "/initial" { next(w, r) return } redirectLogin := "/login?redirect=" + path authBase64, err := r.Cookie("auth") if errors.Is(err, http.ErrNoCookie) { a.entry.Info("No auth provided") http.Redirect(w, r, redirectLogin, http.StatusTemporaryRedirect) return } auth, err := base64.StdEncoding.DecodeString(authBase64.Value) if err != nil { a.entry.Error(err) return } token, err := service.ReadToken(auth, a.key) if err != nil { a.entry.Error(err) http.Redirect(w, r, redirectLogin, http.StatusTemporaryRedirect) return } r = r.WithContext(context.WithValue(r.Context(), "token", token)) a.entry. WithField("userID", token.UserID). WithField("username", token.Username). Info("user recognized") next(w, r) } } func GetTokenFromCtx(w http.ResponseWriter, r *http.Request) *service.Token { tokenValue := r.Context().Value("token") if token, ok := tokenValue.(*service.Token); ok { return token } return nil } type InitialSetupMiddleware struct { userRepository repository.UserRepository } func NewInitialSetupMiddleware(userRepository repository.UserRepository) *InitialSetupMiddleware { return &InitialSetupMiddleware{ userRepository: userRepository, } } func (i *InitialSetupMiddleware) Check(next http.HandlerFunc) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { // if user has been set to context it is logged in already token := GetTokenFromCtx(w, r) if token == nil { next(w, r) return } path := r.URL.Path if path == "/initial" { next(w, r) return } exists, err := i.userRepository.Any(r.Context()) if err != nil { InternalServerError(w, err) return } if exists { next(w, r) return } http.Redirect(w, r, "/initial", http.StatusTemporaryRedirect) } }