internal/rendezvous/rendezvous.go - build - Git at Google

 // Copyright 2023 The Go Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.

 package rendezvous

 import (
 	"context"
 	"fmt"
 	"log"
 	"net"
 	"net/http"
 	"sync"
 	"time"

 	"golang.org/x/build/buildlet"
 	"golang.org/x/build/internal"
 	"golang.org/x/build/revdial/v2"
 	"google.golang.org/api/idtoken"
 )

 // result contains the result for a waiting instance registration.
 type result struct {
 	bc  buildlet.Client
 	err error
 }

 // entry contains the elements needed to process an instance registration.
 type entry struct {
 	deadline time.Time
 	ch       chan *result
 }

 // TokenValidator verifies if a token is valid.
 type TokenValidator func(ctx context.Context, jwt string) bool

 // Rendezvous waits for buildlets to connect, verifies they are valid instances
 // and passes the connection to the waiting caller.
 type Rendezvous struct {
 	mu sync.Mutex

 	m         map[string]*entry
 	validator TokenValidator
 }

 // Option is an optional configuration setting.
 type Option func(*Rendezvous)

 // OptionValidator changes the verifier used by Rendezvous.
 func OptionValidator(v TokenValidator) Option {
 	return func(rdv *Rendezvous) {
 		rdv.validator = v
 	}
 }

 // New creates a Rendezvous element. The context that is passed in should be non-canceled
 // during the lifetime of the running service.
 func New(ctx context.Context, opts ...Option) *Rendezvous {
 	rdv := &Rendezvous{
 		m:         make(map[string]*entry),
 		validator: validateLUCIIDToken,
 	}
 	for _, opt := range opts {
 		opt(rdv)
 	}
 	go internal.PeriodicallyDo(ctx, 10*time.Second, func(ctx context.Context, t time.Time) {
 		rdv.purgeExpiredRegistrations()
 	})
 	return rdv
 }

 // purgeExpiredRegistrations will purge expired registrations.
 func (rdv *Rendezvous) purgeExpiredRegistrations() {
 	rdv.mu.Lock()
 	for id, ent := range rdv.m {
 		if time.Now().After(ent.deadline) {
 			log.Printf("rendezvous: stopped waiting for instance=%q due to timeout", id)
 			ent.ch <- &result{err: fmt.Errorf("timed out waiting for rendezvous client=%q", id)}
 			delete(rdv.m, id)
 		}
 	}
 	rdv.mu.Unlock()
 }

 // RegisterInstance notes an instance and waits for that instance to connect to the handler. An
 // instance must be registered before the instance can attempt to connect. If an instance does
 // not connect before the end of the wait period, the instance will not be able to connect.
 func (rdv *Rendezvous) RegisterInstance(ctx context.Context, id string, wait time.Duration) {
 	rdv.mu.Lock()
 	rdv.m[id] = &entry{
 		deadline: time.Now().Add(wait),
 		ch:       make(chan *result, 1),
 	}
 	rdv.mu.Unlock()
 }

 // DeregisterInstance removes the registration for an instance which has been
 // previously registered.
 func (rdv *Rendezvous) DeregisterInstance(ctx context.Context, id string) {
 	rdv.mu.Lock()
 	delete(rdv.m, id)
 	rdv.mu.Unlock()
 }

 // WaitForInstance waits for the registered instance to successfully connect. It waits for the
 // lifetime of the context. If the instance is not registered or has exceeded the timeout period,
 // it will immediately return an error.
 func (rdv *Rendezvous) WaitForInstance(ctx context.Context, id string) (buildlet.Client, error) {
 	rdv.mu.Lock()
 	e, ok := rdv.m[id]
 	rdv.mu.Unlock()
 	if !ok {
 		return nil, fmt.Errorf("instance not found: name=%q", id)
 	}
 	select {
 	case <-ctx.Done():
 		rdv.mu.Lock()
 		delete(rdv.m, id)
 		rdv.mu.Unlock()
 		return nil, fmt.Errorf("context timeout waiting for rendezvous client=%q: %w", id, ctx.Err())
 	case res := <-e.ch:
 		rdv.mu.Lock()
 		delete(rdv.m, id)
 		close(e.ch)
 		rdv.mu.Unlock()
 		return res.bc, res.err
 	}
 }

 const (
 	// HeaderID is the HTTP header used for passing the gomote ID.
 	HeaderID = "X-Go-Gomote-ID"
 	// HeaderToken is the HTTP header used for passing in the authentication token.
 	HeaderToken = "X-Go-Swarming-Auth-Token"
 	// HeaderHostname is the HTTP header used for passing in the hostname.
 	HeaderHostname = "X-Go-Hostname"
 )

 // HandleReverse handles HTTP requests from the buildlet and passes the connection to
 // the waiter.
 func (rdv *Rendezvous) HandleReverse(w http.ResponseWriter, r *http.Request) {
 	if r.TLS == nil {
 		http.Error(w, "buildlet registration requires SSL", http.StatusInternalServerError)
 		return
 	}
 	var (
 		id        = r.Header.Get(HeaderID)
 		authToken = r.Header.Get(HeaderToken)
 		hostname  = r.Header.Get(HeaderHostname)
 	)
 	if hostname == "" {
 		http.Error(w, "missing X-Go-Hostname header", http.StatusBadRequest)
 		return
 	}
 	if id == "" {
 		http.Error(w, "missing X-Go-Gomote-ID header", http.StatusBadRequest)
 		return
 	}
 	if authToken == "" {
 		http.Error(w, "missing X-Go-Swarming-Auth-Token header", http.StatusBadRequest)
 		return
 	}
 	rdv.mu.Lock()
 	res, ok := rdv.m[id]
 	rdv.mu.Unlock()

 	if !ok {
 		http.Error(w, "not expecting buildlet client", http.StatusPreconditionFailed)
 		return
 	}
 	if !rdv.validator(r.Context(), authToken) {
 		log.Printf("rendezvous: Unable to validate authentication token id=%s", id)
 		http.Error(w, "invalid authentication Token", http.StatusPreconditionFailed)
 		return
 	}
 	hj, ok := w.(http.Hijacker)
 	if !ok {
 		http.Error(w, "webserver does not support hijacking", http.StatusHTTPVersionNotSupported)
 		return
 	}
 	conn, _, err := hj.Hijack()
 	if err != nil {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		res.ch <- &result{err: err}
 		return
 	}
 	bc, err := connToClient(conn, hostname, "swarming_task")
 	if err != nil {
 		log.Printf("rendezvous: unable to create buildlet client: %s", err)
 		conn.Close()
 		res.ch <- &result{err: err}
 		return
 	}
 	res.ch <- &result{bc: bc}
 }

 func connToClient(conn net.Conn, hostname, hostType string) (buildlet.Client, error) {
 	if err := (&http.Response{StatusCode: http.StatusSwitchingProtocols, Proto: "HTTP/1.1"}).Write(conn); err != nil {
 		log.Printf("gomote: error writing upgrade response to reverse buildlet %s (%s) at %s: %v", hostname, hostType, conn.RemoteAddr(), err)
 		conn.Close()
 		return nil, err
 	}
 	revDialer := revdial.NewDialer(conn, "/revdial")
 	revDialerDone := revDialer.Done()
 	dialer := revDialer.Dial

 	client := buildlet.NewClient(conn.RemoteAddr().String(), buildlet.NoKeyPair)
 	client.SetHTTPClient(&http.Client{
 		Transport: &http.Transport{
 			DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
 				return dialer(ctx)
 			},
 		},
 	})
 	client.SetDialer(dialer)
 	client.SetDescription(fmt.Sprintf("reverse peer %s/%s for host type %v", hostname, conn.RemoteAddr(), hostType))

 	var isDead struct {
 		sync.Mutex
 		v bool
 	}
 	client.SetOnHeartbeatFailure(func() {
 		isDead.Lock()
 		isDead.v = true
 		isDead.Unlock()
 		conn.Close()
 	})

 	// If the reverse dialer (which is always reading from the
 	// conn detects that the remote went away, close the buildlet
 	// client proactively.
 	go func() {
 		<-revDialerDone
 		isDead.Lock()
 		defer isDead.Unlock()
 		if !isDead.v {
 			client.Close()
 		}
 	}()
 	tstatus := time.Now()
 	status, err := client.Status(context.Background())
 	if err != nil {
 		log.Printf("Reverse connection %s/%s for %s did not answer status after %v: %v",
 			hostname, conn.RemoteAddr(), hostType, time.Since(tstatus), err)
 		conn.Close()
 		return nil, err
 	}
 	log.Printf("Buildlet %s/%s: %+v for %s", hostname, conn.RemoteAddr(), status, hostType)
 	return client, nil
 }

 // validateLUCIIDToken verifies that the token is valid and contains the expected fields.
 func validateLUCIIDToken(ctx context.Context, jwt string) bool {
 	payload, err := idtoken.Validate(ctx, jwt, "https://gomote.golang.org")
 	if err != nil {
 		log.Printf("rendezvous: unable to validate authentication token: %s", err)
 		return false
 	}
 	if payload.Issuer != "https://accounts.google.com" {
 		log.Printf("rendezvous: incorrect issuer: %q", payload.Issuer)
 		return false
 	}
 	if payload.Expires+30 < time.Now().Unix() || payload.IssuedAt-30 > time.Now().Unix() {
 		log.Printf("rendezvous: Bad JWT times: expires %v, issued %v", time.Unix(payload.Expires, 0), time.Unix(payload.IssuedAt, 0))
 		return false
 	}
 	email, ok := payload.Claims["email"]
 	if !ok || email != "coordinator-builder@golang-ci-luci.iam.gserviceaccount.com" {
 		log.Printf("rendezvous: incorrect email=%s", email)
 		return false
 	}
 	emailVerified, ok := payload.Claims["email_verified"].(bool)
 	if !ok || !emailVerified {
 		log.Printf("rendezvous: email unverified email=%s", email)
 		return false
 	}
 	return true
 }
	// Copyright 2023 The Go Authors. All rights reserved.
	// Use of this source code is governed by a BSD-style
	// license that can be found in the LICENSE file.

	package rendezvous

	import (
	"context"
	"fmt"
	"log"
	"net"
	"net/http"
	"sync"
	"time"

	"golang.org/x/build/buildlet"
	"golang.org/x/build/internal"
	"golang.org/x/build/revdial/v2"
	"google.golang.org/api/idtoken"
	)

	// result contains the result for a waiting instance registration.
	type result struct {
	bc buildlet.Client
	err error
	}

	// entry contains the elements needed to process an instance registration.
	type entry struct {
	deadline time.Time
	ch chan *result
	}

	// TokenValidator verifies if a token is valid.
	type TokenValidator func(ctx context.Context, jwt string) bool

	// Rendezvous waits for buildlets to connect, verifies they are valid instances
	// and passes the connection to the waiting caller.
	type Rendezvous struct {
	mu sync.Mutex

	m map[string]*entry
	validator TokenValidator
	}

	// Option is an optional configuration setting.
	type Option func(*Rendezvous)

	// OptionValidator changes the verifier used by Rendezvous.
	func OptionValidator(v TokenValidator) Option {
	return func(rdv *Rendezvous) {
	rdv.validator = v
	}
	}

	// New creates a Rendezvous element. The context that is passed in should be non-canceled
	// during the lifetime of the running service.
	func New(ctx context.Context, opts ...Option) *Rendezvous {
	rdv := &Rendezvous{
	m: make(map[string]*entry),
	validator: validateLUCIIDToken,
	}
	for _, opt := range opts {
	opt(rdv)
	}
	go internal.PeriodicallyDo(ctx, 10*time.Second, func(ctx context.Context, t time.Time) {
	rdv.purgeExpiredRegistrations()
	})
	return rdv
	}

	// purgeExpiredRegistrations will purge expired registrations.
	func (rdv *Rendezvous) purgeExpiredRegistrations() {
	rdv.mu.Lock()
	for id, ent := range rdv.m {
	if time.Now().After(ent.deadline) {
	log.Printf("rendezvous: stopped waiting for instance=%q due to timeout", id)
	ent.ch <- &result{err: fmt.Errorf("timed out waiting for rendezvous client=%q", id)}
	delete(rdv.m, id)
	}
	}
	rdv.mu.Unlock()
	}

	// RegisterInstance notes an instance and waits for that instance to connect to the handler. An
	// instance must be registered before the instance can attempt to connect. If an instance does
	// not connect before the end of the wait period, the instance will not be able to connect.
	func (rdv *Rendezvous) RegisterInstance(ctx context.Context, id string, wait time.Duration) {
	rdv.mu.Lock()
	rdv.m[id] = &entry{
	deadline: time.Now().Add(wait),
	ch: make(chan *result, 1),
	}
	rdv.mu.Unlock()
	}

	// DeregisterInstance removes the registration for an instance which has been
	// previously registered.
	func (rdv *Rendezvous) DeregisterInstance(ctx context.Context, id string) {
	rdv.mu.Lock()
	delete(rdv.m, id)
	rdv.mu.Unlock()
	}

	// WaitForInstance waits for the registered instance to successfully connect. It waits for the
	// lifetime of the context. If the instance is not registered or has exceeded the timeout period,
	// it will immediately return an error.
	func (rdv *Rendezvous) WaitForInstance(ctx context.Context, id string) (buildlet.Client, error) {
	rdv.mu.Lock()
	e, ok := rdv.m[id]
	rdv.mu.Unlock()
	if !ok {
	return nil, fmt.Errorf("instance not found: name=%q", id)
	}
	select {
	case <-ctx.Done():
	rdv.mu.Lock()
	delete(rdv.m, id)
	rdv.mu.Unlock()
	return nil, fmt.Errorf("context timeout waiting for rendezvous client=%q: %w", id, ctx.Err())
	case res := <-e.ch:
	rdv.mu.Lock()
	delete(rdv.m, id)
	close(e.ch)
	rdv.mu.Unlock()
	return res.bc, res.err
	}
	}

	const (
	// HeaderID is the HTTP header used for passing the gomote ID.
	HeaderID = "X-Go-Gomote-ID"
	// HeaderToken is the HTTP header used for passing in the authentication token.
	HeaderToken = "X-Go-Swarming-Auth-Token"
	// HeaderHostname is the HTTP header used for passing in the hostname.
	HeaderHostname = "X-Go-Hostname"
	)

	// HandleReverse handles HTTP requests from the buildlet and passes the connection to
	// the waiter.
	func (rdv Rendezvous) HandleReverse(w http.ResponseWriter, r http.Request) {
	if r.TLS == nil {
	http.Error(w, "buildlet registration requires SSL", http.StatusInternalServerError)
	return
	}
	var (
	id = r.Header.Get(HeaderID)
	authToken = r.Header.Get(HeaderToken)
	hostname = r.Header.Get(HeaderHostname)
	)
	if hostname == "" {
	http.Error(w, "missing X-Go-Hostname header", http.StatusBadRequest)
	return
	}
	if id == "" {
	http.Error(w, "missing X-Go-Gomote-ID header", http.StatusBadRequest)
	return
	}
	if authToken == "" {
	http.Error(w, "missing X-Go-Swarming-Auth-Token header", http.StatusBadRequest)
	return
	}
	rdv.mu.Lock()
	res, ok := rdv.m[id]
	rdv.mu.Unlock()

	if !ok {
	http.Error(w, "not expecting buildlet client", http.StatusPreconditionFailed)
	return
	}
	if !rdv.validator(r.Context(), authToken) {
	log.Printf("rendezvous: Unable to validate authentication token id=%s", id)
	http.Error(w, "invalid authentication Token", http.StatusPreconditionFailed)
	return
	}
	hj, ok := w.(http.Hijacker)
	if !ok {
	http.Error(w, "webserver does not support hijacking", http.StatusHTTPVersionNotSupported)
	return
	}
	conn, _, err := hj.Hijack()
	if err != nil {
	http.Error(w, err.Error(), http.StatusInternalServerError)
	res.ch <- &result{err: err}
	return
	}
	bc, err := connToClient(conn, hostname, "swarming_task")
	if err != nil {
	log.Printf("rendezvous: unable to create buildlet client: %s", err)
	conn.Close()
	res.ch <- &result{err: err}
	return
	}
	res.ch <- &result{bc: bc}
	}

	func connToClient(conn net.Conn, hostname, hostType string) (buildlet.Client, error) {
	if err := (&http.Response{StatusCode: http.StatusSwitchingProtocols, Proto: "HTTP/1.1"}).Write(conn); err != nil {
	log.Printf("gomote: error writing upgrade response to reverse buildlet %s (%s) at %s: %v", hostname, hostType, conn.RemoteAddr(), err)
	conn.Close()
	return nil, err
	}
	revDialer := revdial.NewDialer(conn, "/revdial")
	revDialerDone := revDialer.Done()
	dialer := revDialer.Dial

	client := buildlet.NewClient(conn.RemoteAddr().String(), buildlet.NoKeyPair)
	client.SetHTTPClient(&http.Client{
	Transport: &http.Transport{
	DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
	return dialer(ctx)
	},
	},
	})
	client.SetDialer(dialer)
	client.SetDescription(fmt.Sprintf("reverse peer %s/%s for host type %v", hostname, conn.RemoteAddr(), hostType))

	var isDead struct {
	sync.Mutex
	v bool
	}
	client.SetOnHeartbeatFailure(func() {
	isDead.Lock()
	isDead.v = true
	isDead.Unlock()
	conn.Close()
	})

	// If the reverse dialer (which is always reading from the
	// conn detects that the remote went away, close the buildlet
	// client proactively.
	go func() {
	<-revDialerDone
	isDead.Lock()
	defer isDead.Unlock()
	if !isDead.v {
	client.Close()
	}
	}()
	tstatus := time.Now()
	status, err := client.Status(context.Background())
	if err != nil {
	log.Printf("Reverse connection %s/%s for %s did not answer status after %v: %v",
	hostname, conn.RemoteAddr(), hostType, time.Since(tstatus), err)
	conn.Close()
	return nil, err
	}
	log.Printf("Buildlet %s/%s: %+v for %s", hostname, conn.RemoteAddr(), status, hostType)
	return client, nil
	}

	// validateLUCIIDToken verifies that the token is valid and contains the expected fields.
	func validateLUCIIDToken(ctx context.Context, jwt string) bool {
	payload, err := idtoken.Validate(ctx, jwt, "https://gomote.golang.org")
	if err != nil {
	log.Printf("rendezvous: unable to validate authentication token: %s", err)
	return false
	}
	if payload.Issuer != "https://accounts.google.com" {
	log.Printf("rendezvous: incorrect issuer: %q", payload.Issuer)
	return false
	}
	if payload.Expires+30 < time.Now().Unix() \|\| payload.IssuedAt-30 > time.Now().Unix() {
	log.Printf("rendezvous: Bad JWT times: expires %v, issued %v", time.Unix(payload.Expires, 0), time.Unix(payload.IssuedAt, 0))
	return false
	}
	email, ok := payload.Claims["email"]
	if !ok \|\| email != "coordinator-builder@golang-ci-luci.iam.gserviceaccount.com" {
	log.Printf("rendezvous: incorrect email=%s", email)
	return false
	}
	emailVerified, ok := payload.Claims["email_verified"].(bool)
	if !ok \|\| !emailVerified {
	log.Printf("rendezvous: email unverified email=%s", email)
	return false
	}
	return true
	}