Server-side token broker

import io.ktor.client.*
import io.ktor.client.call.*
import io.ktor.client.engine.cio.*
import io.ktor.client.plugins.contentnegotiation.*
import io.ktor.client.request.*
import io.ktor.client.request.forms.*
import io.ktor.http.*
import io.ktor.serialization.kotlinx.json.*
import io.ktor.server.application.*
import io.ktor.server.engine.*
import io.ktor.server.netty.*
import io.ktor.server.response.*
import io.ktor.server.routing.*
import kotlinx.serialization.Serializable
import java.time.Instant
import java.util.Base64
import java.util.concurrent.ConcurrentHashMap

private const val AUTH_SERVER_URL = "https://auth.example.com"
private const val CACHE_SAFETY_MARGIN_SECONDS = 60L

private val credentialsByPackage = mapOf(
    "com.example.myapp" to ClientCredentials("android-client-id", "android-client-secret"),
    "com.example.MyApp" to ClientCredentials("ios-client-id", "ios-client-secret"),
    "app.example.com" to ClientCredentials("web-client-id", "web-client-secret"),
)

data class ClientCredentials(val clientId: String, val clientSecret: String)

data class CachedToken(val accessToken: String, val expiresAt: Instant)

@Serializable
data class TokenResponse(val access_token: String, val expires_in: Long)

@Serializable
data class GatekeeperTokenResponse(val access_token: String)

private val tokenCache = ConcurrentHashMap<String, CachedToken>()

private val httpClient = HttpClient(CIO) {
    install(ContentNegotiation) { json() }
}

suspend fun fetchGatekeeperToken(credentials: ClientCredentials): CachedToken {
    val basicAuth = Base64.getEncoder().encodeToString(
        "${credentials.clientId}:${credentials.clientSecret}".toByteArray()
    )

    val response: TokenResponse = httpClient.post("$AUTH_SERVER_URL/oauth2/token") {
        header(HttpHeaders.Authorization, "Basic $basicAuth")
        contentType(ContentType.Application.FormUrlEncoded)
        setBody(FormDataContent(Parameters.build {
            append("grant_type", "client_credentials")
        }))
    }.body()

    val expiresAt = Instant.now().plusSeconds(
        maxOf(response.expires_in - CACHE_SAFETY_MARGIN_SECONDS, 0)
    )

    return CachedToken(response.access_token, expiresAt)
}

suspend fun getCachedToken(applicationPackage: String): String {
    val cached = tokenCache[applicationPackage]
    if (cached != null && cached.expiresAt.isAfter(Instant.now())) {
        return cached.accessToken
    }

    val credentials = credentialsByPackage[applicationPackage]
        ?: error("Unknown application_package: $applicationPackage")

    val freshToken = fetchGatekeeperToken(credentials)
    tokenCache[applicationPackage] = freshToken
    return freshToken.accessToken
}

fun main() {
    embeddedServer(Netty, port = 8080) {
        routing {
            get("/api/gatekeeper-token") {
                val applicationPackage = call.request.queryParameters["application_package"]
                if (applicationPackage.isNullOrBlank()) {
                    call.respond(HttpStatusCode.BadRequest, mapOf("error" to "application_package is required"))
                    return@get
                }

                try {
                    val accessToken = getCachedToken(applicationPackage)
                    call.respond(GatekeeperTokenResponse(accessToken))
                } catch (error: IllegalStateException) {
                    call.respond(HttpStatusCode.BadRequest, mapOf("error" to error.message))
                } catch (error: Exception) {
                    call.respond(HttpStatusCode.BadGateway, mapOf("error" to "Failed to acquire token from Gatekeeper"))
                }
            }
        }
    }.start(wait = true)
}

import org.springframework.http.*;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.client.RestClientException;
import org.springframework.web.client.RestTemplate;

import java.nio.charset.StandardCharsets;
import java.time.Instant;
import java.util.Base64;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@RestController
public class GatekeeperTokenController {

    private static final String AUTH_SERVER_URL = "https://auth.example.com";
    private static final long CACHE_SAFETY_MARGIN_SECONDS = 60;

    private final RestTemplate restTemplate = new RestTemplate();
    private final Map<String, ClientCredentials> credentialsByPackage = Map.of(
            "com.example.myapp", new ClientCredentials("android-client-id", "android-client-secret"),
            "com.example.MyApp", new ClientCredentials("ios-client-id", "ios-client-secret"),
            "app.example.com", new ClientCredentials("web-client-id", "web-client-secret")
    );
    private final Map<String, CachedToken> tokenCache = new ConcurrentHashMap<>();

    @GetMapping("/api/gatekeeper-token")
    public ResponseEntity<?> gatekeeperToken(@RequestParam(value = "application_package", required = false) String applicationPackage) {
        if (applicationPackage == null || applicationPackage.isBlank()) {
            return ResponseEntity.badRequest().body(Map.of("error", "application_package is required"));
        }

        ClientCredentials credentials = credentialsByPackage.get(applicationPackage);
        if (credentials == null) {
            return ResponseEntity.badRequest().body(Map.of("error", "Unknown application_package: " + applicationPackage));
        }

        try {
            String accessToken = getCachedToken(applicationPackage, credentials);
            return ResponseEntity.ok(Map.of("access_token", accessToken));
        } catch (RestClientException ex) {
            return ResponseEntity.status(HttpStatus.BAD_GATEWAY)
                    .body(Map.of("error", "Failed to acquire token from Gatekeeper"));
        }
    }

    private String getCachedToken(String applicationPackage, ClientCredentials credentials) {
        CachedToken cached = tokenCache.get(applicationPackage);
        if (cached != null && cached.expiresAt().isAfter(Instant.now())) {
            return cached.accessToken();
        }

        CachedToken freshToken = fetchGatekeeperToken(credentials);
        tokenCache.put(applicationPackage, freshToken);
        return freshToken.accessToken();
    }

    private CachedToken fetchGatekeeperToken(ClientCredentials credentials) {
        String basicAuth = Base64.getEncoder().encodeToString(
                (credentials.clientId() + ":" + credentials.clientSecret()).getBytes(StandardCharsets.UTF_8)
        );

        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_FORM_URLENCODED);
        headers.set(HttpHeaders.AUTHORIZATION, "Basic " + basicAuth);

        MultiValueMap<String, String> body = new LinkedMultiValueMap<>();
        body.add("grant_type", "client_credentials");

        ResponseEntity<TokenResponse> response = restTemplate.exchange(
                AUTH_SERVER_URL + "/oauth2/token",
                HttpMethod.POST,
                new HttpEntity<>(body, headers),
                TokenResponse.class
        );

        TokenResponse tokenResponse = response.getBody();
        if (tokenResponse == null || tokenResponse.accessToken() == null) {
            throw new RestClientException("Empty token response from Gatekeeper");
        }

        Instant expiresAt = Instant.now().plusSeconds(
                Math.max(tokenResponse.expiresIn() - CACHE_SAFETY_MARGIN_SECONDS, 0)
        );

        return new CachedToken(tokenResponse.accessToken(), expiresAt);
    }

    private record ClientCredentials(String clientId, String clientSecret) {}
    private record CachedToken(String accessToken, Instant expiresAt) {}
    private record TokenResponse(String accessToken, long expiresIn) {}
}

require "sinatra"
require "net/http"
require "json"
require "base64"
require "uri"

AUTH_SERVER_URL = "https://auth.example.com"
CACHE_SAFETY_MARGIN_SECONDS = 60

CREDENTIALS_BY_PACKAGE = {
  "com.example.myapp" => { client_id: "android-client-id", client_secret: "android-client-secret" },
  "com.example.MyApp" => { client_id: "ios-client-id", client_secret: "ios-client-secret" },
  "app.example.com" => { client_id: "web-client-id", client_secret: "web-client-secret" }
}.freeze

$token_cache = {}

def fetch_gatekeeper_token(client_id, client_secret)
  uri = URI("#{AUTH_SERVER_URL}/oauth2/token")
  request = Net::HTTP::Post.new(uri)
  request["Content-Type"] = "application/x-www-form-urlencoded"
  request["Authorization"] = "Basic #{Base64.strict_encode64("#{client_id}:#{client_secret}")}"
  request.body = URI.encode_www_form(grant_type: "client_credentials")

  response = Net::HTTP.start(uri.hostname, uri.port, use_ssl: uri.scheme == "https") do |http|
    http.request(request)
  end

  raise "Gatekeeper auth failed: #{response.code}" unless response.is_a?(Net::HTTPSuccess)

  payload = JSON.parse(response.body)
  expires_at = Time.now.to_i + [payload["expires_in"].to_i - CACHE_SAFETY_MARGIN_SECONDS, 0].max

  {
    access_token: payload["access_token"],
    expires_at: expires_at
  }
end

def get_cached_token(application_package)
  cached = $token_cache[application_package]
  return cached[:access_token] if cached && cached[:expires_at] > Time.now.to_i

  credentials = CREDENTIALS_BY_PACKAGE[application_package]
  raise ArgumentError, "Unknown application_package: #{application_package}" unless credentials

  fresh_token = fetch_gatekeeper_token(credentials[:client_id], credentials[:client_secret])
  $token_cache[application_package] = fresh_token
  fresh_token[:access_token]
end

get "/api/gatekeeper-token" do
  content_type :json

  application_package = params["application_package"]
  halt 400, { error: "application_package is required" }.to_json if application_package.nil? || application_package.empty?

  begin
    access_token = get_cached_token(application_package)
    { access_token: access_token }.to_json
  rescue ArgumentError => e
    halt 400, { error: e.message }.to_json
  rescue StandardError
    halt 502, { error: "Failed to acquire token from Gatekeeper" }.to_json
  end
end

package main

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"
)

const (
	authServerURL           = "https://auth.example.com"
	cacheSafetyMargin       = 60 * time.Second
)

type clientCredentials struct {
	ClientID     string
	ClientSecret string
}

type cachedToken struct {
	AccessToken string
	ExpiresAt   time.Time
}

type tokenResponse struct {
	AccessToken string `json:"access_token"`
	ExpiresIn   int64  `json:"expires_in"`
}

var (
	credentialsByPackage = map[string]clientCredentials{
		"com.example.myapp": {ClientID: "android-client-id", ClientSecret: "android-client-secret"},
		"com.example.MyApp": {ClientID: "ios-client-id", ClientSecret: "ios-client-secret"},
		"app.example.com":   {ClientID: "web-client-id", ClientSecret: "web-client-secret"},
	}

	tokenCache = struct {
		sync.RWMutex
		entries map[string]cachedToken
	}{
		entries: make(map[string]cachedToken),
	}
)

func fetchGatekeeperToken(credentials clientCredentials) (cachedToken, error) {
	form := url.Values{}
	form.Set("grant_type", "client_credentials")

	req, err := http.NewRequest(http.MethodPost, authServerURL+"/oauth2/token", strings.NewReader(form.Encode()))
	if err != nil {
		return cachedToken{}, err
	}

	basicAuth := base64.StdEncoding.EncodeToString(
		[]byte(credentials.ClientID + ":" + credentials.ClientSecret),
	)
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Set("Authorization", "Basic "+basicAuth)

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return cachedToken{}, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return cachedToken{}, fmt.Errorf("gatekeeper auth failed: %s", string(body))
	}

	var payload tokenResponse
	if err := json.NewDecoder(resp.Body).Decode(&payload); err != nil {
		return cachedToken{}, err
	}

	expiresAt := time.Now().Add(time.Duration(payload.ExpiresIn)*time.Second - cacheSafetyMargin)
	if expiresAt.Before(time.Now()) {
		expiresAt = time.Now()
	}

	return cachedToken{
		AccessToken: payload.AccessToken,
		ExpiresAt:   expiresAt,
	}, nil
}

func getCachedToken(applicationPackage string) (string, error) {
	tokenCache.RLock()
	cached, ok := tokenCache.entries[applicationPackage]
	tokenCache.RUnlock()

	if ok && time.Now().Before(cached.ExpiresAt) {
		return cached.AccessToken, nil
	}

	credentials, ok := credentialsByPackage[applicationPackage]
	if !ok {
		return "", fmt.Errorf("unknown application_package: %s", applicationPackage)
	}

	freshToken, err := fetchGatekeeperToken(credentials)
	if err != nil {
		return "", err
	}

	tokenCache.Lock()
	tokenCache.entries[applicationPackage] = freshToken
	tokenCache.Unlock()

	return freshToken.AccessToken, nil
}

func gatekeeperTokenHandler(w http.ResponseWriter, r *http.Request) {
	applicationPackage := r.URL.Query().Get("application_package")
	if applicationPackage == "" {
		writeJSON(w, http.StatusBadRequest, map[string]string{"error": "application_package is required"})
		return
	}

	accessToken, err := getCachedToken(applicationPackage)
	if err != nil {
		if strings.Contains(err.Error(), "unknown application_package") {
			writeJSON(w, http.StatusBadRequest, map[string]string{"error": err.Error()})
			return
		}
		writeJSON(w, http.StatusBadGateway, map[string]string{"error": "Failed to acquire token from Gatekeeper"})
		return
	}

	writeJSON(w, http.StatusOK, map[string]string{"access_token": accessToken})
}

func writeJSON(w http.ResponseWriter, status int, payload any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(payload)
}

func main() {
	http.HandleFunc("/api/gatekeeper-token", gatekeeperTokenHandler)
	if err := http.ListenAndServe(":8080", nil); err != nil && !errors.Is(err, http.ErrServerClosed) {
		panic(err)
	}
}

use axum::{
    extract::Query,
    http::StatusCode,
    response::{IntoResponse, Response},
    routing::get,
    Json, Router,
};
use base64::{engine::general_purpose::STANDARD, Engine};
use reqwest::header::{AUTHORIZATION, CONTENT_TYPE};
use serde::{Deserialize, Serialize};
use std::{
    collections::HashMap,
    sync::{Arc, RwLock},
    time::{Duration, Instant},
};
use tokio::net::TcpListener;

const AUTH_SERVER_URL: &str = "https://auth.example.com";
const CACHE_SAFETY_MARGIN: Duration = Duration::from_secs(60);

#[derive(Clone)]
struct ClientCredentials {
    client_id: String,
    client_secret: String,
}

#[derive(Clone)]
struct CachedToken {
    access_token: String,
    expires_at: Instant,
}

#[derive(Clone)]
struct AppState {
    credentials_by_package: HashMap<String, ClientCredentials>,
    token_cache: Arc<RwLock<HashMap<String, CachedToken>>>,
    http_client: reqwest::Client,
}

#[derive(Deserialize)]
struct TokenQuery {
    application_package: Option<String>,
}

#[derive(Deserialize)]
struct TokenResponse {
    access_token: String,
    expires_in: u64,
}

#[derive(Serialize)]
struct GatekeeperTokenResponse {
    access_token: String,
}

#[derive(Serialize)]
struct ErrorResponse {
    error: String,
}

async fn fetch_gatekeeper_token(
    http_client: &reqwest::Client,
    credentials: &ClientCredentials,
) -> Result<CachedToken, reqwest::Error> {
    let basic_auth = STANDARD.encode(format!(
        "{}:{}",
        credentials.client_id, credentials.client_secret
    ));

    let response = http_client
        .post(format!("{AUTH_SERVER_URL}/oauth2/token"))
        .header(CONTENT_TYPE, "application/x-www-form-urlencoded")
        .header(AUTHORIZATION, format!("Basic {basic_auth}"))
        .body("grant_type=client_credentials")
        .send()
        .await?
        .error_for_status()?
        .json::<TokenResponse>()
        .await?;

    let ttl = Duration::from_secs(response.expires_in).saturating_sub(CACHE_SAFETY_MARGIN);

    Ok(CachedToken {
        access_token: response.access_token,
        expires_at: Instant::now() + ttl,
    })
}

async fn get_cached_token(
    state: &AppState,
    application_package: &str,
) -> Result<String, String> {
    {
        let cache = state.token_cache.read().unwrap();
        if let Some(cached) = cache.get(application_package) {
            if cached.expires_at > Instant::now() {
                return Ok(cached.access_token.clone());
            }
        }
    }

    let credentials = state
        .credentials_by_package
        .get(application_package)
        .cloned()
        .ok_or_else(|| format!("Unknown application_package: {application_package}"))?;

    let fresh_token = fetch_gatekeeper_token(&state.http_client, &credentials)
        .await
        .map_err(|_| "Failed to acquire token from Gatekeeper".to_string())?;

    let access_token = fresh_token.access_token.clone();
    state
        .token_cache
        .write()
        .unwrap()
        .insert(application_package.to_string(), fresh_token);

    Ok(access_token)
}

async fn gatekeeper_token_handler(
    axum::extract::State(state): axum::extract::State<AppState>,
    Query(query): Query<TokenQuery>,
) -> Response {
    let Some(application_package) = query.application_package.filter(|value| !value.is_empty()) else {
        return (
            StatusCode::BAD_REQUEST,
            Json(ErrorResponse {
                error: "application_package is required".to_string(),
            }),
        )
            .into_response();
    };

    match get_cached_token(&state, &application_package).await {
        Ok(access_token) => (
            StatusCode::OK,
            Json(GatekeeperTokenResponse { access_token }),
        )
            .into_response(),
        Err(message) if message.starts_with("Unknown application_package") => (
            StatusCode::BAD_REQUEST,
            Json(ErrorResponse { error: message }),
        )
            .into_response(),
        Err(_) => (
            StatusCode::BAD_GATEWAY,
            Json(ErrorResponse {
                error: "Failed to acquire token from Gatekeeper".to_string(),
            }),
        )
            .into_response(),
    }
}

#[tokio::main]
async fn main() {
    let mut credentials_by_package = HashMap::new();
    credentials_by_package.insert(
        "com.example.myapp".to_string(),
        ClientCredentials {
            client_id: "android-client-id".to_string(),
            client_secret: "android-client-secret".to_string(),
        },
    );
    credentials_by_package.insert(
        "com.example.MyApp".to_string(),
        ClientCredentials {
            client_id: "ios-client-id".to_string(),
            client_secret: "ios-client-secret".to_string(),
        },
    );
    credentials_by_package.insert(
        "app.example.com".to_string(),
        ClientCredentials {
            client_id: "web-client-id".to_string(),
            client_secret: "web-client-secret".to_string(),
        },
    );

    let state = AppState {
        credentials_by_package,
        token_cache: Arc::new(RwLock::new(HashMap::new())),
        http_client: reqwest::Client::new(),
    };

    let app = Router::new()
        .route("/api/gatekeeper-token", get(gatekeeper_token_handler))
        .with_state(state);

    let listener = TcpListener::bind("0.0.0.0:8080").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}