ms-auth/cmd/ms-auth-proxy/main.go

package main

import (
	"context"
	"errors"
	"fmt"
	"git.sch9.ru/new_gate/ms-auth/config"
	delivery "git.sch9.ru/new_gate/ms-auth/internal/users/delivery/grpc"
	userv1gw "git.sch9.ru/new_gate/ms-auth/proto/user/v1"
	"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
	"github.com/ilyakaznacheev/cleanenv"
	"github.com/rs/cors"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/metadata"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"
)

func CustomOutgoingHeaderMatcher(key string) (string, bool) {
	if key == delivery.SessionHeaderName {
		return "Set-Cookie", true
	}

	return fmt.Sprintf("%s%s", runtime.MetadataHeaderPrefix, key), true
}

func CustomIncomingHeaderMatcher(key string) (string, bool) {
	if key == "Cookie" {
		return "Cookie", true
	}

	return fmt.Sprintf("%s%s", runtime.MetadataPrefix, key), true
}

func UnaryClientInterceptor(
	ctx context.Context,
	method string,
	req, reply interface{},
	cc *grpc.ClientConn,
	invoker grpc.UnaryInvoker,
	opts ...grpc.CallOption,
) error {
	md, ok := metadata.FromOutgoingContext(ctx)

	hasCookie := ok && len(md.Get("Cookie")) > 0

	if hasCookie {
		cookies, err := http.ParseCookie(md.Get("Cookie")[0])
		if err != nil {
			return err
		}

		if len(cookies) != 1 {
			return errors.New("invalid cookie")
		}

		md.Set(delivery.SessionHeaderName, cookies[0].Value)
	}

	ctx = metadata.NewOutgoingContext(ctx, md)

	err := invoker(ctx, method, req, reply, cc, opts...)
	if err != nil {
		return err
	}

	for _, o := range opts {
		header, ok := o.(grpc.HeaderCallOption)
		if !ok {
			continue
		}

		values := header.HeaderAddr.Get(delivery.SessionHeaderName)

		if len(values) != 1 {
			continue
		}

		sessionId := values[0]

		cookie := http.Cookie{
			Name:     "SESSIONID",
			Value:    sessionId,
			Path:     "/",
			HttpOnly: true,
		}

		if len(sessionId) == 0 {
			cookie.Expires = time.Unix(0, 0)
		} else {
			cookie.MaxAge = 3600
		}

		header.HeaderAddr.Set(delivery.SessionHeaderName, cookie.String())
	}

	return nil
}

func main() {
	var cfg config.Config
	err := cleanenv.ReadConfig(".env", &cfg)
	if err != nil {
		panic(fmt.Sprintf("error reading config: %s", err.Error()))
	}

	ctx := context.Background()
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	mux := runtime.NewServeMux(runtime.WithOutgoingHeaderMatcher(CustomOutgoingHeaderMatcher), runtime.WithIncomingHeaderMatcher(CustomIncomingHeaderMatcher))
	opts := []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithChainUnaryInterceptor(UnaryClientInterceptor)}
	err = userv1gw.RegisterUserServiceHandlerFromEndpoint(ctx, mux, cfg.Address, opts)
	if err != nil {
		panic(err)
	}

	c := cors.New(cors.Options{
		AllowedOrigins:   []string{"http://*", "https://*"},
		AllowedMethods:   []string{"GET", "POST", "PUT", "DELETE", "OPTIONS"},
		AllowedHeaders:   []string{"Content-Type", "Set-Cookie", "Credentials"},
		ExposedHeaders:   []string{"Link"},
		AllowCredentials: true,
	})

	go func() {
		err = http.ListenAndServe(cfg.ProxyAddress, c.Handler(mux))
		if err != nil {
			panic(err)
		}
	}()

	fmt.Println("server proxy started")

	stop := make(chan os.Signal, 1)
	signal.Notify(stop, syscall.SIGTERM, syscall.SIGINT)

	<-stop
	return
}