| package main |
| |
| import ( |
| "context" |
| "flag" |
| "fmt" |
| "net/http" |
| "os" |
| |
| "github.com/gorilla/mux" |
| "go.skia.org/infra/go/common" |
| "go.skia.org/infra/go/httputils" |
| "go.skia.org/infra/go/sklog" |
| "go.skia.org/infra/skolo/go/service_accounts" |
| "go.skia.org/infra/skolo/go/skmetadata" |
| ) |
| |
| const ( |
| // APP_NAME is the name of this app. |
| APP_NAME = "metadata_server" |
| ) |
| |
| var ( |
| // Flags. |
| port = flag.String("port", ":8000", "HTTP service port for the web server (e.g., ':8000')") |
| promPort = flag.String("prom_port", ":20000", "Metrics service address (e.g., ':10110')") |
| |
| serviceAccountIpMapping = map[*service_accounts.ServiceAccount][]string{ |
| service_accounts.ChromeSwarming: {"*"}, |
| service_accounts.ChromiumSwarm: {"*"}, |
| service_accounts.Jumphost: {"self"}, |
| service_accounts.RpiMaster: {"192.168.1.98", "192.168.1.99"}, |
| } |
| ) |
| |
| type instanceMetadataMap map[string]map[string]string |
| |
| func (m instanceMetadataMap) Get(instance, key string) (string, error) { |
| rv, ok := m[instance][key] |
| if !ok { |
| return "", fmt.Errorf("Unknown instance or key.") |
| } |
| return rv, nil |
| } |
| |
| type projectMetadataMap map[string]string |
| |
| func (m projectMetadataMap) Get(key string) (string, error) { |
| rv, ok := m[key] |
| if !ok { |
| return "", fmt.Errorf("Unknown key.") |
| } |
| return rv, nil |
| } |
| |
| func main() { |
| common.InitWithMust( |
| APP_NAME, |
| common.PrometheusOpt(promPort), |
| //common.CloudLoggingOpt(), |
| ) |
| |
| // TODO(borenet): Load these from a file? |
| var pm projectMetadataMap = map[string]string{ |
| "mykey": "myvalue", |
| } |
| var im instanceMetadataMap = map[string]map[string]string{ |
| "inst": { |
| "mykey2": "myvalue2", |
| }, |
| } |
| |
| hostname, err := os.Hostname() |
| if err != nil { |
| sklog.Fatal(err) |
| } |
| serviceAccounts, ok := service_accounts.JumphostServiceAccountMapping[hostname] |
| if !ok { |
| sklog.Fatalf("Hostname not in jumphost mapping: %s", hostname) |
| } |
| |
| tokens := make(map[string]*skmetadata.ServiceAccountToken, len(serviceAccounts)) |
| tokenIpMapping := make(map[string]string, len(serviceAccounts)) |
| for _, acct := range serviceAccounts { |
| ipAddrs, ok := serviceAccountIpMapping[acct] |
| if !ok { |
| sklog.Fatalf("Service account has no IP address mapping: %s", acct.Email) |
| } |
| tokenFile := fmt.Sprintf("/var/local/token_%s.json", acct.Nickname) |
| |
| for _, ipAddr := range ipAddrs { |
| tokenIpMapping[ipAddr] = tokenFile |
| if _, ok := tokens[tokenFile]; !ok { |
| tok, err := skmetadata.NewServiceAccountToken(tokenFile, false) |
| if err != nil { |
| sklog.Fatal(err) |
| } |
| go tok.UpdateLoop(context.Background()) |
| tokens[tokenFile] = tok |
| } |
| } |
| } |
| tokenMapping := make(map[string]*skmetadata.ServiceAccountToken, len(tokenIpMapping)) |
| for ipAddr, tokenFile := range tokenIpMapping { |
| tokenMapping[ipAddr] = tokens[tokenFile] |
| } |
| |
| r := mux.NewRouter() |
| skmetadata.SetupServer(r, pm, im, tokenMapping) |
| http.Handle("/", httputils.LoggingGzipRequestResponse(r)) |
| sklog.Infof("Ready to serve on http://localhost%s", *port) |
| sklog.Fatal(http.ListenAndServe(*port, nil)) |
| } |