github.com/grafana/synthetic-monitoring-agent@v0.64.2/cmd/test-api/internal/grpc/grpc.go (about)

     1  package grpc
     2  
     3  import (
     4  	"context"
     5  	"encoding/base64"
     6  	"errors"
     7  	"fmt"
     8  	"math"
     9  	"net"
    10  	"strings"
    11  	"time"
    12  
    13  	"github.com/gogo/status"
    14  	sm "github.com/grafana/synthetic-monitoring-agent/pkg/pb/synthetic_monitoring"
    15  	"github.com/rs/zerolog"
    16  	"google.golang.org/grpc"
    17  	"google.golang.org/grpc/codes"
    18  	"google.golang.org/grpc/keepalive"
    19  	"google.golang.org/grpc/metadata"
    20  )
    21  
    22  type ServerDb interface {
    23  	FindProbeIDByToken(ctx context.Context, token []byte) (int64, error)
    24  }
    25  
    26  type Opts struct {
    27  	Logger            zerolog.Logger
    28  	ListenAddr        string
    29  	ChecksServer      sm.ChecksServer
    30  	TenantsServer     sm.TenantsServer
    31  	AdHocChecksServer sm.AdHocChecksServer
    32  	Db                ServerDb
    33  }
    34  
    35  type Server struct {
    36  	srv    *grpc.Server
    37  	logger zerolog.Logger
    38  	addr   string
    39  	db     ServerDb
    40  }
    41  
    42  func NewServer(ctx context.Context, opts *Opts) (*Server, error) {
    43  	srv := &Server{
    44  		logger: opts.Logger,
    45  		addr:   opts.ListenAddr,
    46  		db:     opts.Db,
    47  	}
    48  
    49  	srvOpts := []grpc.ServerOption{
    50  		grpc.ChainUnaryInterceptor(srv.unaryHandler),
    51  		grpc.ChainStreamInterceptor(srv.streamHandler),
    52  		grpc.KeepaliveParams(keepalive.ServerParameters{
    53  			MaxConnectionIdle:     5 * time.Minute,
    54  			MaxConnectionAge:      time.Duration(math.MaxInt64),
    55  			MaxConnectionAgeGrace: time.Duration(math.MaxInt64),
    56  			Time:                  120 * time.Second,
    57  			Timeout:               20 * time.Second,
    58  		}),
    59  	}
    60  
    61  	srv.srv = grpc.NewServer(srvOpts...)
    62  
    63  	if opts.ChecksServer != nil {
    64  		sm.RegisterChecksServer(srv.srv, opts.ChecksServer)
    65  	}
    66  
    67  	if opts.TenantsServer != nil {
    68  		sm.RegisterTenantsServer(srv.srv, opts.TenantsServer)
    69  	}
    70  
    71  	if opts.AdHocChecksServer != nil {
    72  		sm.RegisterAdHocChecksServer(srv.srv, opts.AdHocChecksServer)
    73  	}
    74  
    75  	return srv, nil
    76  }
    77  
    78  func (s *Server) Run(ctx context.Context) error {
    79  	s.logger.Info().Str("address", s.addr).Msg("starting GRPC server")
    80  
    81  	l, err := (&net.ListenConfig{}).Listen(ctx, "tcp", s.addr)
    82  	if err != nil {
    83  		return err
    84  	}
    85  
    86  	if err := s.srv.Serve(l); err != nil {
    87  		return fmt.Errorf("could not serve GRPC on %s: %w", s.addr, err)
    88  	}
    89  
    90  	return nil
    91  }
    92  
    93  func (s *Server) GracefulStop() {
    94  	s.srv.GracefulStop()
    95  }
    96  
    97  func (s *Server) Stop() {
    98  	s.srv.Stop()
    99  }
   100  
   101  func (s *Server) unaryHandler(ctx context.Context, req any, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
   102  	probeID, err := s.validateToken(ctx)
   103  	if err != nil {
   104  		return nil, err
   105  	}
   106  
   107  	return handler(contextWithProbeId(ctx, probeID), req)
   108  }
   109  
   110  func (s *Server) streamHandler(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
   111  	probeID, err := s.validateToken(ss.Context())
   112  	if err != nil {
   113  		return err
   114  	}
   115  
   116  	wrapper := &serverStreamWrapper{
   117  		ctx:          contextWithProbeId(ss.Context(), probeID),
   118  		ServerStream: ss,
   119  	}
   120  
   121  	return handler(srv, wrapper)
   122  }
   123  
   124  func (s *Server) validateToken(ctx context.Context) (int64, error) {
   125  	md, ok := metadata.FromIncomingContext(ctx)
   126  	if !ok {
   127  		return -1, errors.New("missing grpc metadata")
   128  	}
   129  
   130  	auth, found := md["authorization"]
   131  	if !found {
   132  		msg := "missing authorization token"
   133  		err := status.Error(codes.InvalidArgument, msg)
   134  
   135  		return -1, err
   136  	}
   137  
   138  	if len(auth) != 1 {
   139  		msg := "invalid authorization metadata"
   140  		err := status.Error(codes.InvalidArgument, msg)
   141  
   142  		return -1, err
   143  	}
   144  
   145  	const prefix = "Bearer "
   146  
   147  	if !strings.HasPrefix(auth[0], prefix) {
   148  		msg := "invalid authorization metadata format"
   149  		err := status.Error(codes.InvalidArgument, msg)
   150  
   151  		return -1, err
   152  	}
   153  
   154  	b64token := strings.TrimPrefix(auth[0], prefix)
   155  	if _, err := base64.StdEncoding.DecodeString(b64token); err != nil {
   156  		msg := "invalid authorization encoding"
   157  		err := status.Error(codes.InvalidArgument, msg)
   158  
   159  		return -1, err
   160  	}
   161  
   162  	id, err := s.db.FindProbeIDByToken(ctx, []byte(b64token))
   163  	if err != nil {
   164  		return -1, err
   165  	}
   166  
   167  	return id, nil
   168  }
   169  
   170  type probeIdKey struct{}
   171  
   172  func contextWithProbeId(ctx context.Context, probeId int64) context.Context {
   173  	return context.WithValue(ctx, probeIdKey{}, probeId)
   174  }
   175  
   176  func probeIdFromContext(ctx context.Context) (int64, bool) {
   177  	id, ok := ctx.Value(probeIdKey{}).(int64)
   178  	return id, ok
   179  }
   180  
   181  type serverStreamWrapper struct {
   182  	ctx context.Context
   183  	grpc.ServerStream
   184  }
   185  
   186  func (w *serverStreamWrapper) Context() context.Context {
   187  	return w.ctx
   188  }