github.com/stevenmatthewt/agent@v3.5.4+incompatible/agent/api_proxy.go (about)

     1  package agent
     2  
     3  import (
     4  	"crypto/sha256"
     5  	"fmt"
     6  	"io/ioutil"
     7  	"net"
     8  	"net/http"
     9  	"net/http/httputil"
    10  	"net/url"
    11  	"os"
    12  	"runtime"
    13  	"sync"
    14  	"time"
    15  
    16  	"github.com/buildkite/agent/api"
    17  	"github.com/buildkite/agent/logger"
    18  )
    19  
    20  // APIProxy provides either a unix socket or a tcp socket listener with a proxy
    21  // that will authenticate to the Buildkite Agent API
    22  type APIProxy struct {
    23  	upstreamToken    string
    24  	upstreamEndpoint string
    25  	token            string
    26  	socket           *os.File
    27  	listener         net.Listener
    28  	listenerWg       *sync.WaitGroup
    29  }
    30  
    31  func NewAPIProxy(endpoint string, token string) *APIProxy {
    32  	var wg sync.WaitGroup
    33  	wg.Add(1)
    34  
    35  	return &APIProxy{
    36  		upstreamToken:    token,
    37  		upstreamEndpoint: endpoint,
    38  		token:            fmt.Sprintf("%x", sha256.Sum256([]byte(string(time.Now().UnixNano())))),
    39  		listenerWg:       &wg,
    40  	}
    41  }
    42  
    43  // Listen on either a tcp socket (for windows) or a unix socket
    44  func (p *APIProxy) Listen() error {
    45  	defer p.listenerWg.Done()
    46  	var err error
    47  
    48  	// windows doesn't support unix sockets, so we fall back to a tcp socket
    49  	if runtime.GOOS == `windows` {
    50  		p.listener, err = p.listenOnTCPSocket()
    51  	} else {
    52  		p.listener, p.socket, err = p.listenOnUnixSocket()
    53  	}
    54  
    55  	if err != nil {
    56  		return err
    57  	}
    58  
    59  	endpoint, err := url.Parse(p.upstreamEndpoint)
    60  	if err != nil {
    61  		return err
    62  	}
    63  
    64  	go func() {
    65  		proxy := httputil.NewSingleHostReverseProxy(endpoint)
    66  		proxy.Transport = &api.AuthenticatedTransport{Token: p.upstreamToken}
    67  
    68  		// customize the reverse proxy director so that we can make some changes to the request
    69  		director := proxy.Director
    70  		proxy.Director = func(req *http.Request) {
    71  			director(req)
    72  
    73  			// set the host header whilst proxying
    74  			req.Host = req.URL.Host
    75  		}
    76  
    77  		// serve traffic, proxy off to the reverse proxy
    78  		_ = http.Serve(p.listener, http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
    79  			if r.Header.Get(`Authorization`) != `Token `+p.token {
    80  				http.Error(rw, "Invalid authorization token", http.StatusBadRequest)
    81  				return
    82  			}
    83  			proxy.ServeHTTP(rw, r)
    84  		}))
    85  	}()
    86  
    87  	return nil
    88  }
    89  
    90  func (p *APIProxy) listenOnUnixSocket() (net.Listener, *os.File, error) {
    91  	socket, err := ioutil.TempFile("", "agent-socket")
    92  	if err != nil {
    93  		return nil, nil, err
    94  	}
    95  
    96  	// Servers should unlink the socket path name prior to binding it.
    97  	// https://troydhanson.github.io/network/Unix_domain_sockets.html
    98  	_ = os.Remove(socket.Name())
    99  
   100  	logger.Debug("[APIProxy] Listening on unix socket %s", socket.Name())
   101  
   102  	// create a unix socket to do the listening
   103  	l, err := net.Listen("unix", socket.Name())
   104  	if err != nil {
   105  		return nil, nil, err
   106  	}
   107  
   108  	// Restrict to owner r+w permissions
   109  	if err = os.Chmod(socket.Name(), 0600); err != nil {
   110  		return nil, nil, err
   111  	}
   112  
   113  	return l, socket, nil
   114  }
   115  
   116  func (p *APIProxy) listenOnTCPSocket() (net.Listener, error) {
   117  	// Listen on the first available non-privileged port
   118  	l, err := net.Listen("tcp", "127.0.0.1:0")
   119  	if err != nil {
   120  		return nil, err
   121  	}
   122  
   123  	logger.Debug("[APIProxy] Listening on tcp socket %s", l.Addr().String())
   124  	return l, nil
   125  }
   126  
   127  // Close any listeners or internal files
   128  func (p *APIProxy) Close() error {
   129  	defer p.listenerWg.Add(1)
   130  	return p.listener.Close()
   131  }
   132  
   133  // Wait blocks until the listener is ready
   134  func (p *APIProxy) Wait() {
   135  	p.listenerWg.Wait()
   136  }
   137  
   138  func (p *APIProxy) Endpoint() string {
   139  	if p.socket != nil {
   140  		return `unix://` + p.listener.Addr().String()
   141  	}
   142  	return `http://` + p.listener.Addr().String()
   143  }
   144  
   145  func (p *APIProxy) AccessToken() string {
   146  	return p.token
   147  }