package main

// This program emulates the meta data server in GCE and is intended to be used
// in the Skolo just for serving tokens to swarming clients. Note that the
// serving path is hard-coded to
// /computeMetadata/v1/instance/service-accounts/default/token.
//
// Note that the token comes from using Google default credentials, so the
// GOOGLE_APPLICATION_CREDENTIALS environment variable must be set.

import (
	"encoding/json"
	"flag"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/mux"
	"go.skia.org/infra/go/auth"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/httputils"
	"go.skia.org/infra/go/metrics2"
	"go.skia.org/infra/go/skerr"
	"go.skia.org/infra/go/sklog"
	"golang.org/x/oauth2"
)

// refreshInterval is how often we should refresh the token.
var refreshInterval = 10 * time.Minute

// Flags.
var (
	local    = flag.Bool("local", false, "True if running locally, not in prod.")
	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')")
)

type server struct {
	successfulRefresh metrics2.Counter
	failedRefresh     metrics2.Counter

	// mutex protects tokenSource and latestToken.
	mutex       sync.Mutex
	tokenSource oauth2.TokenSource
	latestToken *oauth2.Token
}

// newServer creates a new *server with a running Go routine that refreshes the token.
func newServer() (*server, error) {
	tokenSource, err := auth.NewDefaultTokenSource(*local, auth.SCOPE_USERINFO_EMAIL)
	if err != nil {
		return nil, skerr.Wrapf(err, "Failed to create token source.")
	}

	ret := &server{
		tokenSource:       tokenSource,
		successfulRefresh: metrics2.GetCounter("metadata_successful_refresh"),
		failedRefresh:     metrics2.GetCounter("metadata_failed_refresh"),
	}

	if err := ret.step(); err != nil {
		return nil, skerr.Wrapf(err, "Failed to get token.")
	}

	go ret.refresh()

	return ret, nil
}

func (s *server) handleTokenRequest(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	s.mutex.Lock()
	t := s.latestToken
	s.mutex.Unlock()
	// Copied from
	// https://github.com/golang/oauth2/blob/f6093e37b6cb4092101a298aba5d794eb570757f/google/google.go#L185
	res := struct {
		AccessToken  string `json:"access_token"`
		ExpiresInSec int    `json:"expires_in"`
		TokenType    string `json:"token_type"`
	}{
		AccessToken:  t.AccessToken,
		ExpiresInSec: int(t.Expiry.Sub(time.Now()).Seconds()),
		TokenType:    t.TokenType,
	}
	sklog.Infof("Token requested by %s, serving %s", r.RemoteAddr, res.AccessToken[len(res.AccessToken)-8:])
	if err := json.NewEncoder(w).Encode(res); err != nil {
		httputils.ReportError(w, err, "Failed to write response.", http.StatusInternalServerError)
		return
	}
}

// step does a single refresh of the oauth token.
func (s *server) step() error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	// We need to re-create the tokenSource to force it to fetch a fresh token.
	// See https://github.com/golang/oauth2/pull/396.
	tokenSource, err := auth.NewDefaultTokenSource(*local, auth.SCOPE_USERINFO_EMAIL)
	if err != nil {
		return skerr.Wrapf(err, "Failed to create token source.")
	}
	s.tokenSource = tokenSource

	latestToken, err := s.tokenSource.Token()
	if err != nil {
		s.failedRefresh.Inc(1)
		return skerr.Wrapf(err, "Failed to retrieve token.")
	}

	sklog.Info("Successfully refreshed.")
	s.successfulRefresh.Inc(1)
	s.latestToken = latestToken
	return nil
}

// refresh handles periodically refreshing the token.
//
// This function does not return and should be run as a Go routine.
func (s *server) refresh() {
	for range time.Tick(refreshInterval) {
		if err := s.step(); err != nil {
			sklog.Errorf("Failed to refresh token: %s")
		}
	}
}

func main() {
	common.InitWithMust(
		"metadata",
		common.PrometheusOpt(promPort),
	)

	s, err := newServer()
	if err != nil {
		sklog.Fatal(err)
	}

	r := mux.NewRouter()
	r.HandleFunc("/computeMetadata/v1/instance/service-accounts/default/token", s.handleTokenRequest)
	http.Handle("/", httputils.Healthz(httputils.LoggingGzipRequestResponse(r)))
	sklog.Infof("Ready to serve on http://localhost%s", *port)
	sklog.Fatal(http.ListenAndServe(*port, nil))
}