purposelimiter.go

// This file contains the interceptor function that is used to
// perform purpose limiting data minimization operations within
// a servide-side gRPC response interceptor.
//
//
// The interceptor function is called by the gRPC server like this:
// path to the JWT's public key file is passed as a parameter keyPath
// grpc.UnaryInterceptor(purposelimiter.UnaryServerInterceptor(keyPath))

package naive_approach

import (
	"context"
	"crypto/rsa"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"github.com/golang-jwt/jwt/v5"
	"io/ioutil"
	"log"
	"math"
	"math/rand"
	"strconv"
	"strings"
	"time"

	_ "github.com/Siar-Akbayin/jwt-go-auth"
	"github.com/google/differential-privacy/go/dpagg"
	"github.com/google/differential-privacy/go/noise"
	"google.golang.org/grpc"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/reflect/protoreflect"
)

// CustomClaims is our custom metadata
type CustomClaims struct {
	Policy struct {
		Allowed     map[string][]string `json:"allowed"`
		Generalized map[string][]string `json:"generalized"`
		Noised      map[string][]string `json:"noised"`
		Reduced     map[string][]string `json:"reduced"`
	} `json:"policy"`

	jwt.RegisteredClaims
}

func UnaryServerInterceptor(policyPath string, serviceName string, purpose string, privateKey string, publicKey string) grpc.UnaryServerInterceptor {
	return interceptor(policyPath, serviceName, purpose, privateKey, publicKey)
}

func interceptor(policyPath string, serviceName string, purpose string, privateKey string, publicKey string) grpc.UnaryServerInterceptor {
	return func(
		ctx context.Context,
		req interface{},
		info *grpc.UnaryServerInfo,
		handler grpc.UnaryHandler,
	) (interface{}, error) {

		h, err := handler(ctx, req)
		if err != nil {
			return nil, err
		}

		publicKey, err := loadPublicKey(publicKey)
		if err != nil {
			return nil, err
		}

		// get token
		token, err := generateToken(policyPath, serviceName, purpose, privateKey)
		if err != nil {
			return nil, err
		}
		tkn, err := jwt.ParseWithClaims(token, &CustomClaims{}, func(token *jwt.Token) (interface{}, error) {
			return publicKey, nil
		})
		// ----------------------
		// !	Validation		!
		// ----------------------

		//if err != nil {
		//	return nil, err
		//}
		//
		//if !tkn.Valid {
		//	return nil, fmt.Errorf("invalid token")
		//}

		claims := tkn.Claims.(*CustomClaims)
		// ----------------------
		// !	Validation		!
		// ----------------------

		// Check if the response is a proto.Message
		msg, ok := h.(proto.Message)
		if !ok {
			return nil, fmt.Errorf("response is not a proto.Message")
		}

		// Invoke ProtoReflect() to get a protoreflect.Message
		reflectedMsg := msg.ProtoReflect()
		// Declare a slice to store field names
		var fieldNames []string

		reflectedMsg.Range(func(fd protoreflect.FieldDescriptor, v protoreflect.Value) bool {
			name := fd.TextName()

			fieldNames = append(fieldNames, name)

			return true
		})
		// Iterate over the fields of the message
		for _, field := range fieldNames {
			// Check if the field is not in the allowed list
			// --> Pass if the field is allowed
			if !contains(claims.Policy.Allowed, field) {
				// Check if the field is in one of the minimized lists
				if contains(claims.Policy.Generalized, field) {
					// Generalize the field
					switch reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field)).Kind() {
					case protoreflect.Int32Kind:
						reflectedMsg.Set(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field)), protoreflect.ValueOf(generalizeIntParam(reflectedMsg.Get(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field))).Int(), claims.Policy.Generalized[field][1])))
					case protoreflect.StringKind:
						reflectedMsg.Set(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field)), protoreflect.ValueOf(generalizeStringParam(reflectedMsg.Get(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field))).String(), claims.Policy.Generalized[field][1])))
					case protoreflect.FloatKind:
						reflectedMsg.Set(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field)), protoreflect.ValueOf(generalizeFloatParam(reflectedMsg.Get(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field))).Float(), claims.Policy.Generalized[field][1])))
					}
				} else if contains(claims.Policy.Noised, field) {
					// Noise the field
					switch reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field)).Kind() {
					case protoreflect.Int32Kind:
						reflectedMsg.Set(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field)), protoreflect.ValueOf(noiseIntParam(reflectedMsg.Get(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field))).Int(), claims.Policy.Noised[field][1])))
					case protoreflect.StringKind:
						reflectedMsg.Set(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field)), protoreflect.ValueOf(noiseStringParam(reflectedMsg.Get(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field))).String(), claims.Policy.Noised[field][1])))
					case protoreflect.FloatKind:
						reflectedMsg.Set(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field)), protoreflect.ValueOf(noiseFloatParam(reflectedMsg.Get(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field))).Float(), claims.Policy.Noised[field][1])))
					}
				} else if contains(claims.Policy.Reduced, field) {
					// Reduce the field
					switch reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field)).Kind() {
					case protoreflect.Int32Kind:
						reflectedMsg.Set(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field)), protoreflect.ValueOf(reduceIntParam(reflectedMsg.Get(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field))).Int(), claims.Policy.Reduced[field][1])))
					case protoreflect.StringKind:
						reflectedMsg.Set(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field)), protoreflect.ValueOf(reduceStringParam(reflectedMsg.Get(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field))).String(), claims.Policy.Reduced[field][1])))
					case protoreflect.FloatKind:
						reflectedMsg.Set(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field)), protoreflect.ValueOf(reduceFloatParam(reflectedMsg.Get(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field))).Float(), claims.Policy.Reduced[field][1])))
					}
				} else {
					//Suppress the field
					switch reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field)).Kind() {
					case protoreflect.Int32Kind:
						reflectedMsg.Set(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field)), protoreflect.ValueOf(suppressInt(reflectedMsg.Get(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field))).Int())))
					case protoreflect.StringKind:
						reflectedMsg.Set(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field)), protoreflect.ValueOf(suppressString(reflectedMsg.Get(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field))).String())))
					case protoreflect.FloatKind:
						reflectedMsg.Set(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field)), protoreflect.ValueOf(suppressFloat(reflectedMsg.Get(reflectedMsg.Descriptor().Fields().ByName(protoreflect.Name(field))).Float())))
					}
				}
			}
		}
		return h, nil
	}
}

func generateToken(policyPath string, serviceName string, purpose string, privateKey string) (string, error) {

	// Load policy from file
	policyData, err := ioutil.ReadFile(policyPath)
	if err != nil {
		log.Fatalf("Error reading policy.json: %v", err)
	}

	// Parse the policy JSON into a map
	var policyMap map[string]interface{}
	err = json.Unmarshal(policyData, &policyMap)
	if err != nil {
		log.Fatalf("Error parsing policy.json: %v", err)
	}

	// Retrieve the services object from the policy
	servicesObj, exists := policyMap["services"].(map[string]interface{})
	if !exists {
		return "", fmt.Errorf("Invalid policy format: services not found")
	}

	// Get the service policy based on the service name
	servicePolicy, exists := servicesObj[serviceName].(map[string]interface{})
	if !exists {
		return "", fmt.Errorf("Service %s not found in policy file", serviceName)
	}

	// Get the policy for the specified purpose
	purposePolicy, exists := servicePolicy[purpose].(map[string]interface{})
	if !exists {
		return "", fmt.Errorf("Purpose %s not found in service policy", purpose)
	}

	// Create the reduced policy based on the purpose policy
	reducedPolicy := map[string]interface{}{
		"allowed":     purposePolicy["allowed"],
		"generalized": purposePolicy["generalized"],
		"noised":      purposePolicy["noised"],
		"reduced":     purposePolicy["reduced"],
	}

	// Convert the reduced policy to JSON
	reducedPolicyJSON, err := json.Marshal(reducedPolicy)
	if err != nil {
		log.Fatalf("Error marshaling reduced policy: %v", err)
	}

	// Load the RSA private key from file
	keyData, err := ioutil.ReadFile(privateKey)
	if err != nil {
		log.Fatalf("Error reading private key: %v", err)
	}

	// Parse the RSA private key
	privateKeyy, err := jwt.ParseRSAPrivateKeyFromPEM(keyData)
	if err != nil {
		log.Fatalf("Error parsing private key: %v", err)
	}

	// Create the Claims
	claims := struct {
		Policy json.RawMessage `json:"policy"`
		jwt.RegisteredClaims
	}{
		reducedPolicyJSON,
		jwt.RegisteredClaims{
			// Valid for 2 hrs
			ExpiresAt: jwt.NewNumericDate(time.Now().Add(time.Hour * 2)),
			Issuer:    "tokenGenerator",
		},
	}

	token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
	tokenString, err := token.SignedString(privateKeyy)
	if err != nil {
		log.Fatalf("Error signing token: %v", err)
		return "", err
	}
	return tokenString, nil
}

// ------ minimzation functions ------

// Suppression functions
func suppressInt(number int64) int32 {
	// receives an integer (e.g., house number) and returns -1 as "none".
	return -1
}
func suppressFloat(number float64) float64 {
	// receives a float (e.g., house number) and returns -1 as "none".
	return -1
}
func suppressString(text string) string {
	// receives a string (e.g., street name) and cuts it off after the 5th character.
	return ""
}

// Noising functions
// --> parametrized
func noiseIntParam(number int64, param string) int64 {
	// receives an int and returns noised version of it.
	// rand.Int63n returns a non-negative pseudo-random 63-bit integer
	// available noise functions:
	// - Gaussian
	// - Laplace
	// e.g. noiseIntParam(135, "Gaussian")

	// Gaussian noise
	var n noise.Noise
	epsilon := 1.0
	delta := 0.0

	switch param {
	case "Gaussian":
		delta = 0.01
		n = noise.Gaussian()
	case "Laplace":
		n = noise.Laplace()
	default:
		log.Fatalf("Error: Unknown noise function: %v", param)
		return -1
	}

	// Instantiate a new BoundedSum with the chosen noise mechanism.
	sumParams := &dpagg.BoundedSumInt64Options{
		Epsilon:                  epsilon,
		Delta:                    delta,
		Noise:                    n,
		MaxPartitionsContributed: 1,
		Lower:                    0,
		Upper:                    100,
	}
	sum := dpagg.NewBoundedSumInt64(sumParams)

	// Add our number to the sum
	sum.Add(number)

	// Calculate the result with noise
	result := sum.Result()

	// The result is a float64, so we'll convert it to int64
	return int64(math.Abs(float64(result)))
}
func noiseFloatParam(number float64, param string) float64 {
	// receives an int and returns noised version of it.
	// rand.Int63n returns a non-negative pseudo-random 63-bit integer
	// available noise functions:
	// - Gaussian
	// - Laplace
	// e.g. noiseIntParam(135, "Gaussian")

	// Gaussian noise
	var n noise.Noise
	epsilon := 1.0
	delta := 0.0

	switch param {
	case "Gaussian":
		delta = 0.01
		n = noise.Gaussian()
	case "Laplace":
		n = noise.Laplace()
	default:
		log.Fatalf("Error: Unknown noise function: %v", param)
		return -1
	}

	// Instantiate a new BoundedSum with the chosen noise mechanism.
	sumParams := &dpagg.BoundedSumFloat64Options{
		Epsilon:                  epsilon,
		Delta:                    delta,
		Noise:                    n,
		MaxPartitionsContributed: 1,
		Lower:                    0,
		Upper:                    100,
	}
	sum := dpagg.NewBoundedSumFloat64(sumParams)

	// Add our number to the sum
	sum.Add(number)

	// Calculate the result with noise
	result := sum.Result()

	// The result is a float64, so we'll convert it to int64
	return math.Abs(float64(result))
}
func noiseStringParam(string, param string) string {
	// currently not implemented
	// suppressing the field instead
	return suppressString(string)
}

// --> non-parametrized
func noiseInt(number int64) int64 {
	// receives a house number and returns noised version of it.
	// rand.Int31 returns a non-negative pseudo-random 31-bit integer as an int32 from the default Source.
	return number - rand.Int63n(number) + rand.Int63n(number)
}
func noiseString(string) string {
	// receives a string and returns noised version of it.
	return ""
}

// Generalization functions
// --> parametrized
func generalizeIntParam(number int64, param string) int64 {
	// receives an integer (e.g., house number) and returns its range of param's as the lower end of the interval specified by param..
	// e.g. generalizeIntParam(135, 10) -> 131

	intParam, err := strconv.ParseInt(param, 10, 64)
	if err != nil {
		log.Fatalf("Error on converting string to int: %v", err)
	}
	return number/intParam*intParam + 1
}
func generalizeFloatParam(number float64, param string) float64 {
	// receives a float (e.g., house number) and returns its range of 10's as the lower end of the interval.
	// e.g. 135.0 -> 131.0
	floatParam, err := strconv.ParseFloat(param, 64)
	if err != nil {
		log.Fatalf("Error on converting string to float: %v", err)
	}
	return number/floatParam*floatParam + 1
}
func generalizeStringParam(text string, param string) string {
	// receives a string (e.g., street name) and returns the first ncharacter(s), with n=param.
	intParam, err := strconv.ParseInt(param, 10, 64)
	if err != nil {
		log.Fatalf("Error on converting string to int: %v", err)
	}
	return text[0:intParam]
}

// --> non-parametrized
func generalizeInt(number int64) int64 {
	// receives an integer (e.g., house number) and returns its range of 10's as the lower end of the interval.
	// e.g. 135 -> 131
	return number/10*10 + 1
}
func generalizeString(text string) string {
	// receives a string (e.g., street name) and returns the first character.
	return text[0:1]
}

// Reduction functions
// --> parametrized
func reduceIntParam(number int64, param string) int64 {
	// receives an integer and divides it by the specified by param.
	// e.g. reduceIntParam(135, 10) -> 13,5

	intParam, err := strconv.ParseInt(param, 10, 64)
	if err != nil {
		log.Fatalf("Error on converting string to int: %v", err)
	}
	return number / intParam * intParam
}
func reduceFloatParam(number float64, param string) float64 {
	// receives a float and divides it by the specified by param.
	// e.g. reduceFloatParam(135.0, 10) -> 13.5

	floatParam, err := strconv.ParseFloat(param, 64)
	if err != nil {
		log.Fatalf("Error on converting string to float: %v", err)
	}
	return float64(number) / floatParam * floatParam
}
func reduceStringParam(text string, param string) string {
	// receives a string (e.g., street name) and returns the first n character(s), with n=param.
	intParam, err := strconv.ParseInt(param, 10, 64)
	if err != nil {
		log.Fatalf("Error on converting string to int: %v", err)
	}
	return text[0:intParam]
}

// --> non-parametrized
func reduceInt(number int64) int64 {
	return number / 10
}
func reduceString(text string) string {
	// receives a string (e.g., street name) and returns the first 4 characters.
	return text[0:3]
}

// ------ utiliy functions ------

// contains checks if a field is present in a map
func contains(m map[string][]string, key string) bool {
	_, ok := m[key]
	return ok
}

// getLastPart returns the last part of a string separated by dots
// e.g., main.HelloReply.name --> name
func getLastPart(s string) (string, error) {
	parts := strings.Split(s, ".")
	if len(parts) < 1 {
		return "", errors.New("input string is empty")
	}
	return parts[len(parts)-1], nil
}

// loadPublicKey loads a public key from a file
func loadPublicKey(path string) (*rsa.PublicKey, error) {
	pubPEMData, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}

	block, _ := pem.Decode(pubPEMData)
	if block == nil || block.Type != "PUBLIC KEY" {
		return nil, errors.New("failed to decode PEM block containing public key")
	}

	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, err
	}

	rsaPub, ok := pub.(*rsa.PublicKey)
	if !ok {
		return nil, errors.New("public key is not in RSA format")
	}

	return rsaPub, nil
}