baron-sso/backend/internal/service/keto_service.go

package service

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"os"
	"time"
)

type KetoService interface {
	CheckPermission(ctx context.Context, subject, namespace, object, relation string) (bool, error)
	CreateRelation(ctx context.Context, namespace, object, relation, subject string) error
	DeleteRelation(ctx context.Context, namespace, object, relation, subject string) error
	ListRelations(ctx context.Context, namespace, object, relation, subject string) ([]RelationTuple, error)
	ListObjects(ctx context.Context, namespace, relation, subject string) ([]string, error)
}

type ketoService struct {
	readURL  string
	writeURL string
	client   *http.Client
}

func NewKetoService() KetoService {
	readURL := os.Getenv("KETO_READ_URL")
	if readURL == "" {
		readURL = "http://keto:4466"
	}
	writeURL := os.Getenv("KETO_WRITE_URL")
	if writeURL == "" {
		writeURL = "http://keto:4467"
	}

	return &ketoService{
		readURL:  readURL,
		writeURL: writeURL,
		client:   &http.Client{},
	}
}

type RelationTuple struct {
	Namespace string `json:"namespace"`
	Object    string `json:"object"`
	Relation  string `json:"relation"`
	SubjectID string `json:"subject_id"`
}

type relationTuplesResponse struct {
	RelationTuples []RelationTuple `json:"relation_tuples"`
	NextPageToken  string          `json:"next_page_token"`
}

func (s *ketoService) ListRelations(ctx context.Context, namespace, object, relation, subject string) ([]RelationTuple, error) {
	u, _ := url.Parse(fmt.Sprintf("%s/relation-tuples", s.readURL))
	q := u.Query()
	if namespace != "" {
		q.Set("namespace", namespace)
	}
	if object != "" {
		q.Set("object", object)
	}
	if relation != "" {
		q.Set("relation", relation)
	}
	if subject != "" {
		q.Set("subject_id", subject)
	}
	u.RawQuery = q.Encode()

	req, _ := http.NewRequestWithContext(ctx, "GET", u.String(), nil)
	resp, err := s.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("keto returned status %d: %s", resp.StatusCode, string(body))
	}

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

	return res.RelationTuples, nil
}

type checkResponse struct {
	Allowed bool `json:"allowed"`
}

func (s *ketoService) CheckPermission(ctx context.Context, subject, namespace, object, relation string) (bool, error) {
	u, _ := url.Parse(fmt.Sprintf("%s/relation-tuples/check", s.readURL))
	q := u.Query()
	q.Set("namespace", namespace)
	q.Set("object", object)
	q.Set("relation", relation)
	q.Set("subject_id", subject)
	u.RawQuery = q.Encode()

	var lastErr error
	maxRetries := 5
	backoff := 200 * time.Millisecond

	for i := 0; i < maxRetries; i++ {
		req, _ := http.NewRequestWithContext(ctx, "GET", u.String(), nil)
		resp, err := s.client.Do(req)
		if err == nil {
			defer resp.Body.Close()
			if resp.StatusCode == http.StatusOK {
				var res checkResponse
				if err := json.NewDecoder(resp.Body).Decode(&res); err != nil {
					return false, err
				}
				return res.Allowed, nil
			}
			if resp.StatusCode == http.StatusForbidden {
				return false, nil
			}
			body, _ := io.ReadAll(resp.Body)
			lastErr = fmt.Errorf("keto returned status %d: %s", resp.StatusCode, string(body))
		} else {
			lastErr = err
		}

		if i < maxRetries-1 {
			slog.Debug("Retrying Keto CheckPermission...", "attempt", i+1, "error", lastErr)
			time.Sleep(backoff)
			backoff *= 2
		}
	}

	return false, lastErr
}

func (s *ketoService) CreateRelation(ctx context.Context, namespace, object, relation, subject string) error {
	u := fmt.Sprintf("%s/admin/relation-tuples", s.writeURL)
	payload := map[string]interface{}{
		"namespace":  namespace,
		"object":     object,
		"relation":   relation,
		"subject_id": subject,
	}
	body, _ := json.Marshal(payload)

	// Exponential Backoff Retry Logic
	var lastErr error
	maxRetries := 5
	backoff := 200 * time.Millisecond

	for i := 0; i < maxRetries; i++ {
		req, _ := http.NewRequestWithContext(ctx, "PUT", u, bytes.NewReader(body))
		req.Header.Set("Content-Type", "application/json")

		resp, err := s.client.Do(req)
		if err == nil {
			defer resp.Body.Close()
			if resp.StatusCode == http.StatusCreated || resp.StatusCode == http.StatusNoContent || resp.StatusCode == http.StatusOK {
				slog.Debug("Keto relation created", "namespace", namespace, "object", object, "relation", relation, "subject", subject)
				return nil
			}
			resBody, _ := io.ReadAll(resp.Body)
			lastErr = fmt.Errorf("keto returned status %d: %s", resp.StatusCode, string(resBody))
		} else {
			lastErr = err
		}

		if i < maxRetries-1 {
			slog.Debug("Retrying Keto CreateRelation...", "attempt", i+1, "error", lastErr)
			time.Sleep(backoff)
			backoff *= 2
		}
	}

	slog.Error("Keto create relation failed after retries", "error", lastErr, "namespace", namespace, "object", object, "relation", relation, "subject", subject)
	return lastErr
}

func (s *ketoService) DeleteRelation(ctx context.Context, namespace, object, relation, subject string) error {
	u, _ := url.Parse(fmt.Sprintf("%s/admin/relation-tuples", s.writeURL))
	q := u.Query()
	q.Set("namespace", namespace)
	q.Set("object", object)
	q.Set("relation", relation)
	q.Set("subject_id", subject)
	u.RawQuery = q.Encode()

	var lastErr error
	maxRetries := 5
	backoff := 200 * time.Millisecond

	for i := 0; i < maxRetries; i++ {
		req, _ := http.NewRequestWithContext(ctx, "DELETE", u.String(), nil)
		resp, err := s.client.Do(req)
		if err == nil {
			defer resp.Body.Close()
			if resp.StatusCode == http.StatusNoContent || resp.StatusCode == http.StatusOK {
				slog.Debug("Keto relation deleted", "namespace", namespace, "object", object, "relation", relation, "subject", subject)
				return nil
			}
			resBody, _ := io.ReadAll(resp.Body)
			lastErr = fmt.Errorf("keto returned status %d: %s", resp.StatusCode, string(resBody))
		} else {
			lastErr = err
		}

		if i < maxRetries-1 {
			slog.Debug("Retrying Keto DeleteRelation...", "attempt", i+1, "error", lastErr)
			time.Sleep(backoff)
			backoff *= 2
		}
	}

	slog.Error("Keto delete relation failed after retries", "error", lastErr)
	return lastErr
}

func (s *ketoService) ListObjects(ctx context.Context, namespace, relation, subject string) ([]string, error) {
	u, _ := url.Parse(fmt.Sprintf("%s/relation-tuples", s.readURL))
	q := u.Query()
	if namespace != "" {
		q.Set("namespace", namespace)
	}
	if relation != "" {
		q.Set("relation", relation)
	}
	if subject != "" {
		q.Set("subject_id", subject)
	}
	u.RawQuery = q.Encode()

	req, _ := http.NewRequestWithContext(ctx, "GET", u.String(), nil)
	resp, err := s.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("keto returned status %d: %s", resp.StatusCode, string(body))
	}

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

	objects := make([]string, 0, len(res.RelationTuples))
	seen := make(map[string]bool)
	for _, rt := range res.RelationTuples {
		if !seen[rt.Object] {
			objects = append(objects, rt.Object)
			seen[rt.Object] = true
		}
	}

	return objects, nil
}