Add http to https redirect trust center

Signed-off-by: Bryan Frimin <bryan@getprobo.com>

pkg/server/server.go

@@ -177,24 +177,6 @@ func (s *Server) setExtraHeaders(w http.ResponseWriter) {
 }
 
 func (s *Server) handleCustomDomain404(w http.ResponseWriter, r *http.Request) {
-	ctx := r.Context()
-
-	if r.TLS == nil {
-		domain := r.Host
-
-		_, err := s.proboService.LoadOrganizationByDomain(ctx, domain)
-		if err == nil {
-			httpsURL := "https://" + r.Host + r.URL.RequestURI()
-			s.logger.InfoCtx(ctx, "404 on HTTP custom domain, redirecting to HTTPS",
-				log.String("domain", domain),
-				log.String("from", r.URL.RequestURI()),
-				log.String("to", httpsURL),
-			)
-			http.Redirect(w, r, httpsURL, http.StatusMovedPermanently)
-			return
-		}
-	}
-
 	httpserver.RenderError(w, http.StatusNotFound, errors.New("not found"))
 }
 
@@ -262,13 +244,19 @@ func (s *Server) loadTrustCenterByDomain(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		ctx := r.Context()
 
-		if r.TLS == nil || r.TLS.ServerName == "" {
+		// For HTTP requests, use r.Host; for HTTPS requests, use r.TLS.ServerName
+		var domain string
+		if r.TLS != nil && r.TLS.ServerName != "" {
+			domain = r.TLS.ServerName
+		} else {
+			domain = r.Host
+		}
+
+		if domain == "" {
 			next.ServeHTTP(w, r)
 			return
 		}
 
-		domain := r.TLS.ServerName
-
 		s.logger.InfoCtx(ctx, "loading organization by custom domain",
 			log.String("domain", domain),
 			log.String("path", r.URL.Path),
@@ -319,6 +307,26 @@ func (s *Server) stripTrustPrefix(next http.Handler) http.Handler {
 	})
 }
 
+func (s *Server) redirectHTTPToHTTPSForCustomDomain(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		ctx := r.Context()
+
+		if r.TLS == nil {
+			if ctx.Value(trust_v1.CustomDomainOrganizationIDKey) != nil {
+				httpsURL := "https://" + r.Host + r.URL.RequestURI()
+				s.logger.InfoCtx(ctx, "HTTP request to custom domain, redirecting to HTTPS",
+					log.String("domain", r.Host),
+					log.String("path", r.URL.Path),
+				)
+				http.Redirect(w, r, httpsURL, http.StatusMovedPermanently)
+				return
+			}
+		}
+
+		next.ServeHTTP(w, r)
+	})
+}
+
 func (s *Server) trustCenterRouter() chi.Router {
 	r := chi.NewRouter()
 
@@ -340,6 +348,7 @@ func (s *Server) TrustCenterHandler() http.Handler {
 	})
 
 	r.Use(s.loadTrustCenterByDomain)
+	r.Use(s.redirectHTTPToHTTPSForCustomDomain)
 	r.NotFound(s.handleCustomDomain404)
 
 	r.Mount("/", s.trustCenterRouter())