github.com/m-lab/locate@v0.17.6/cmd/monitoring-token/monitoring-token.go (about)

     1  package main
     2  
     3  import (
     4  	"context"
     5  	"flag"
     6  	"log"
     7  	"net/http"
     8  	"net/url"
     9  	"os"
    10  	"os/exec"
    11  	"time"
    12  
    13  	"github.com/m-lab/access/token"
    14  	"github.com/m-lab/go/flagx"
    15  	"github.com/m-lab/go/logx"
    16  	"github.com/m-lab/go/pretty"
    17  	"github.com/m-lab/go/rtx"
    18  	v2 "github.com/m-lab/locate/api/v2"
    19  	"github.com/m-lab/locate/proxy"
    20  	"github.com/m-lab/locate/static"
    21  	"gopkg.in/square/go-jose.v2/jwt"
    22  )
    23  
    24  var (
    25  	locate    = flagx.MustNewURL("http://localhost:8080/v2/platform/monitoring/")
    26  	privKey   flagx.FileBytes
    27  	machine   string
    28  	service   string
    29  	timeout   time.Duration
    30  	envName   string
    31  	envValue  string
    32  	logFatalf = log.Fatalf
    33  	// TODO (kinkade): Remove these two variables and corresponding flag
    34  	// definitions once all monitoring clients are migrated to use only
    35  	// monitoring URLs.
    36  	serviceURL        bool
    37  	serviceURLKeyName string
    38  )
    39  
    40  func init() {
    41  	setupFlags()
    42  }
    43  
    44  func setupFlags() {
    45  	flag.Var(&locate, "locate-url", "URL Prefix for locate service")
    46  	flag.Var(&privKey, "monitoring-signer-key", "Private JWT key used for signing")
    47  	flag.StringVar(&machine, "machine", "", "Machine name used as Audience in the jwt Claim")
    48  	flag.StringVar(&service, "service", "ndt/ndt5", "<experiment>/<datatype> to request monitoring access tokens")
    49  	flag.DurationVar(&timeout, "timeout", 60*time.Second, "Complete request and command execution within timeout")
    50  	flag.StringVar(&envName, "env-name", "MONITORING_URL", "Export the monitoring locate URL to the named environment variable before executing given command")
    51  	flag.BoolVar(&serviceURL, "service-url", false, "Return a service URL instead of the default monitoring locate URL")
    52  	flag.StringVar(&serviceURLKeyName, "service-url-key-name", "wss://:3010/ndt_protocol", "The key name to extract from the monitoring locate result Target.URLs")
    53  }
    54  
    55  func main() {
    56  	flag.Parse()
    57  	rtx.Must(flagx.ArgsFromEnvWithLog(flag.CommandLine, false), "Failed to read args from env")
    58  
    59  	// This process signs access tokens for /v2/platform/monitoring requests to the
    60  	// locate service. NOTE: the locate service MUST be configured with the
    61  	// corresponding public key to verify these access tokens.
    62  	priv, err := token.NewSigner(privKey)
    63  	rtx.Must(err, "Failed to allocate signer")
    64  
    65  	// Create a claim, similar to the locate service, and sign it.
    66  	cl := jwt.Claims{
    67  		Issuer:   static.IssuerMonitoring,
    68  		Subject:  machine,
    69  		Audience: jwt.Audience{static.AudienceLocate},
    70  		Expiry:   jwt.NewNumericDate(time.Now().Add(time.Minute)),
    71  	}
    72  
    73  	// Signing the claim generates the access token string.
    74  	logx.Debug.Println(cl)
    75  	token, err := priv.Sign(cl)
    76  	rtx.Must(err, "Failed to sign claims")
    77  
    78  	// Add the token to the URL parameters in the request to the locate service.
    79  	params, err := url.ParseQuery(locate.RawQuery)
    80  	rtx.Must(err, "failed to parse given query")
    81  	params.Set("access_token", token)
    82  	locate.RawQuery = params.Encode()
    83  	locate.Path = locate.Path + service
    84  
    85  	// Prepare a context with absolute timeout for getting token and running command.
    86  	ctx, cancel := context.WithTimeout(context.Background(), timeout)
    87  	defer cancel()
    88  
    89  	// TODO (kinkade): Once all monitoring clients are migrated to use only
    90  	// monitoring URLs then this whole conditional block can be removed and
    91  	// envValue will always just contain the monitoring URL.
    92  	if serviceURL {
    93  		logx.Debug.Println("Issue request to:", locate.URL)
    94  		req, err := http.NewRequestWithContext(ctx, http.MethodGet, locate.String(), nil)
    95  		rtx.Must(err, "Failed to create request from url: %q", locate.URL)
    96  
    97  		// Get monitoring result.
    98  		mr := &v2.MonitoringResult{}
    99  		_, err = proxy.UnmarshalResponse(req, mr)
   100  		rtx.Must(err, "Failed to get response")
   101  		logx.Debug.Println(pretty.Sprint(mr))
   102  		if mr.Error != nil {
   103  			logFatalf("ERROR: %s %s", mr.Error.Title, mr.Error.Detail)
   104  			return
   105  		}
   106  		if mr.Target == nil {
   107  			logFatalf("ERROR: monitoring result Target field is nil!")
   108  			return
   109  		}
   110  		envValue = mr.Target.URLs[serviceURLKeyName]
   111  	} else {
   112  		envValue = locate.String()
   113  	}
   114  
   115  	// Place the URL into the named environment variable for access by the command.
   116  	os.Setenv(envName, envValue)
   117  	logx.Debug.Println("Setting:", envName, "=", envValue)
   118  	logx.Debug.Println("Exec:", flag.Args())
   119  	args := flag.Args()
   120  	if len(args) == 0 {
   121  		logFatalf("ERROR: no command given to execute")
   122  		return
   123  	}
   124  	cmd := exec.CommandContext(ctx, args[0], args[1:]...)
   125  	if logx.LogxDebug.Get() {
   126  		cmd.Stdout = os.Stdout
   127  		cmd.Stderr = os.Stderr
   128  	}
   129  	rtx.Must(cmd.Run(), "Failed to run %#v", args)
   130  }