diff --git a/management/internals/server/server_multiplex_test.go b/management/internals/server/server_multiplex_test.go new file mode 100644 index 000000000..2b87e974c --- /dev/null +++ b/management/internals/server/server_multiplex_test.go @@ -0,0 +1,109 @@ +package server + +import ( + "context" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" + "fmt" + "io" + "math/big" + "net" + "net/http" + "testing" + "time" + + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials" + "google.golang.org/grpc/health" + healthpb "google.golang.org/grpc/health/grpc_health_v1" +) + +func newSelfSignedCert(t *testing.T) tls.Certificate { + t.Helper() + + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + require.NoError(t, err) + + template := x509.Certificate{ + SerialNumber: big.NewInt(1), + Subject: pkix.Name{CommonName: "127.0.0.1"}, + NotBefore: time.Now().Add(-time.Hour), + NotAfter: time.Now().Add(time.Hour), + IPAddresses: []net.IP{net.ParseIP("127.0.0.1")}, + } + der, err := x509.CreateCertificate(rand.Reader, &template, &template, &key.PublicKey, key) + require.NoError(t, err) + + return tls.Certificate{Certificate: [][]byte{der}, PrivateKey: key} +} + +func TestServeMultiplexedRoutesProtocols(t *testing.T) { + tcpListener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + t.Cleanup(func() { _ = tcpListener.Close() }) + + baseTLSConfig := &tls.Config{ + Certificates: []tls.Certificate{newSelfSignedCert(t)}, + NextProtos: []string{"h2", "http/1.1"}, + } + tlsListener := tls.NewListener(tcpListener, preferHTTP1ForDualProtoClients(baseTLSConfig)) + + grpcServer := grpc.NewServer() + healthpb.RegisterHealthServer(grpcServer, health.NewServer()) + t.Cleanup(grpcServer.Stop) + + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = fmt.Fprintf(w, "proto=%d", r.ProtoMajor) + }) + + s := &BaseServer{errCh: make(chan error, 4)} + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + s.serveMultiplexed(ctx, tlsListener, grpcServer, handler, true) + + addr := tcpListener.Addr().String() + url := "https://" + addr + "/" + + grpcConn, err := grpc.NewClient(addr, + grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{InsecureSkipVerify: true}))) + require.NoError(t, err) + t.Cleanup(func() { _ = grpcConn.Close() }) + + checkCtx, checkCancel := context.WithTimeout(ctx, 5*time.Second) + defer checkCancel() + resp, err := healthpb.NewHealthClient(grpcConn).Check(checkCtx, &healthpb.HealthCheckRequest{}) + require.NoError(t, err) + require.Equal(t, healthpb.HealthCheckResponse_SERVING, resp.Status) + + get := func(client *http.Client) string { + t.Helper() + res, err := client.Get(url) + require.NoError(t, err) + body, err := io.ReadAll(res.Body) + require.NoError(t, err) + _ = res.Body.Close() + return string(body) + } + + dualProtoClient := &http.Client{ + Timeout: 5 * time.Second, + Transport: &http.Transport{ + TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, + ForceAttemptHTTP2: true, + }, + } + require.Equal(t, "proto=1", get(dualProtoClient), "dual-ALPN client should be steered to HTTP/1.1") + + h1OnlyClient := &http.Client{ + Timeout: 5 * time.Second, + Transport: &http.Transport{ + TLSClientConfig: &tls.Config{InsecureSkipVerify: true, NextProtos: []string{"http/1.1"}}, + }, + } + require.Equal(t, "proto=1", get(h1OnlyClient)) +} diff --git a/management/internals/shared/grpc/token_refresh_scheduler.go b/management/internals/shared/grpc/token_refresh_scheduler.go new file mode 100644 index 000000000..dbf5d73b1 --- /dev/null +++ b/management/internals/shared/grpc/token_refresh_scheduler.go @@ -0,0 +1,152 @@ +package grpc + +import ( + "container/heap" + "context" + "sync" + "time" +) + +const ( + refreshWorkerCount = 4 + refreshWorkQueueSize = 1024 +) + +type refreshKind int + +const ( + refreshKindTURN refreshKind = iota + refreshKindRelay +) + +type refreshJob struct { + ctx context.Context + accountID string + peerID string + kind refreshKind + interval time.Duration + nextRun time.Time + index int +} + +type refreshJobHeap []*refreshJob + +func (h refreshJobHeap) Len() int { return len(h) } +func (h refreshJobHeap) Less(i, j int) bool { return h[i].nextRun.Before(h[j].nextRun) } + +func (h refreshJobHeap) Swap(i, j int) { + h[i], h[j] = h[j], h[i] + h[i].index = i + h[j].index = j +} + +func (h *refreshJobHeap) Push(x any) { + job := x.(*refreshJob) + job.index = len(*h) + *h = append(*h, job) +} + +func (h *refreshJobHeap) Pop() any { + old := *h + n := len(old) + job := old[n-1] + old[n-1] = nil + job.index = -1 + *h = old[:n-1] + return job +} + +// refreshScheduler executes periodic credential refresh jobs for all peers +// from one timer goroutine and a fixed worker pool, instead of two parked +// goroutines per connected peer. +type refreshScheduler struct { + mu sync.Mutex + jobs refreshJobHeap + wake chan struct{} + work chan *refreshJob + run func(job *refreshJob) +} + +func newRefreshScheduler(run func(job *refreshJob)) *refreshScheduler { + s := &refreshScheduler{ + wake: make(chan struct{}, 1), + work: make(chan *refreshJob, refreshWorkQueueSize), + run: run, + } + go s.loop() + for range refreshWorkerCount { + go s.worker() + } + return s +} + +func (s *refreshScheduler) schedule(job *refreshJob) { + s.mu.Lock() + job.nextRun = time.Now().Add(job.interval) + heap.Push(&s.jobs, job) + s.mu.Unlock() + + select { + case s.wake <- struct{}{}: + default: + } +} + +func (s *refreshScheduler) cancel(job *refreshJob) { + s.mu.Lock() + defer s.mu.Unlock() + if job.index >= 0 { + heap.Remove(&s.jobs, job.index) + } +} + +func (s *refreshScheduler) loop() { + timer := time.NewTimer(time.Hour) + if !timer.Stop() { + <-timer.C + } + + for { + s.mu.Lock() + now := time.Now() + var due []*refreshJob + for len(s.jobs) > 0 && !s.jobs[0].nextRun.After(now) { + job := s.jobs[0] + job.nextRun = job.nextRun.Add(job.interval) + if !job.nextRun.After(now) { + job.nextRun = now.Add(job.interval) + } + heap.Fix(&s.jobs, 0) + due = append(due, job) + } + wait := time.Duration(-1) + if len(s.jobs) > 0 { + wait = time.Until(s.jobs[0].nextRun) + } + s.mu.Unlock() + + for _, job := range due { + s.work <- job + } + + if wait < 0 { + <-s.wake + continue + } + + timer.Reset(wait) + select { + case <-s.wake: + if !timer.Stop() { + <-timer.C + } + case <-timer.C: + } + } +} + +func (s *refreshScheduler) worker() { + for job := range s.work { + s.run(job) + } +}