From 3cdcf025e3bb18a0164a1dbff0f81ce7fa0205d2 Mon Sep 17 00:00:00 2001
From: Pavel Ven Gulbin <2tvenom@gmail.com>
Date: Fri, 1 Dec 2023 18:40:33 +0300
Subject: [PATCH] Added sign func for external signing method

- Added default signing method for local key
- Added `WithCustomSignMethod` builder option
---
 biscuit.go      | 22 ++++++++++++++--------
 biscuit_test.go |  2 +-
 builder.go      | 25 ++++++++++++++++++++-----
 3 files changed, 35 insertions(+), 14 deletions(-)

diff --git a/biscuit.go b/biscuit.go
index ecc5bd7..f202942 100644
--- a/biscuit.go
+++ b/biscuit.go
@@ -2,20 +2,17 @@ package biscuit
 
 import (
 	"bytes"
+	"crypto/ed25519"
 	"crypto/rand"
 	"encoding/binary"
-
-	//"crypto/sha256"
-	"crypto/ed25519"
 	"errors"
 	"fmt"
 	"io"
 
+	"google.golang.org/protobuf/proto"
+
 	"github.com/biscuit-auth/biscuit-go/v2/datalog"
 	"github.com/biscuit-auth/biscuit-go/v2/pb"
-
-	//"github.com/biscuit-auth/biscuit-go/sig"
-	"google.golang.org/protobuf/proto"
 )
 
 // Biscuit represents a valid Biscuit token
@@ -53,7 +50,12 @@ var (
 	UnsupportedAlgorithm = errors.New("biscuit: unsupported signature algorithm")
 )
 
-func New(rng io.Reader, root ed25519.PrivateKey, baseSymbols *datalog.SymbolTable, authority *Block) (*Biscuit, error) {
+func New(
+	rng io.Reader,
+	sign signFunc,
+	baseSymbols *datalog.SymbolTable,
+	authority *Block,
+) (*Biscuit, error) {
 	if rng == nil {
 		rng = rand.Reader
 	}
@@ -83,7 +85,11 @@ func New(rng io.Reader, root ed25519.PrivateKey, baseSymbols *datalog.SymbolTabl
 	toSign := append(marshalledAuthority[:], toSignAlgorithm...)
 	toSign = append(toSign, nextPublicKey[:]...)
 
-	signature := ed25519.Sign(root, toSign)
+	var signature []byte
+	if signature, err = sign(toSign); err != nil {
+		return nil, err
+	}
+
 	nextKey := &pb.PublicKey{
 		Algorithm: &algorithm,
 		Key:       nextPublicKey,
diff --git a/biscuit_test.go b/biscuit_test.go
index 667e2f5..b1e7b58 100644
--- a/biscuit_test.go
+++ b/biscuit_test.go
@@ -402,7 +402,7 @@ func TestNewErrors(t *testing.T) {
 
 	t.Run("authority block Strings overlap", func(t *testing.T) {
 		_, privateRoot, _ := ed25519.GenerateKey(rng)
-		_, err := New(rng, privateRoot, &datalog.SymbolTable{"String1", "String2"}, &Block{
+		_, err := New(rng, DefaultSignMethod(privateRoot), &datalog.SymbolTable{"String1", "String2"}, &Block{
 			symbols: &datalog.SymbolTable{"String1"},
 		})
 		require.Equal(t, ErrSymbolTableOverlap, err)
diff --git a/builder.go b/builder.go
index 599fa68..f1bba5d 100644
--- a/builder.go
+++ b/builder.go
@@ -27,8 +27,8 @@ type Builder interface {
 }
 
 type builder struct {
-	rng  io.Reader
-	root ed25519.PrivateKey
+	rng        io.Reader
+	signMethod signFunc
 
 	symbolsStart int
 	symbols      *datalog.SymbolTable
@@ -38,7 +38,16 @@ type builder struct {
 	context      string
 }
 
-type builderOption func(b *builder)
+type (
+	builderOption func(b *builder)
+	signFunc      func([]byte) ([]byte, error)
+)
+
+func DefaultSignMethod(root ed25519.PrivateKey) signFunc {
+	return func(toSign []byte) ([]byte, error) {
+		return ed25519.Sign(root, toSign), nil
+	}
+}
 
 func WithRandom(rng io.Reader) builderOption {
 	return func(b *builder) {
@@ -53,13 +62,19 @@ func WithSymbols(symbols *datalog.SymbolTable) builderOption {
 	}
 }
 
+func WithCustomSignMethod(s signFunc) builderOption {
+	return func(b *builder) {
+		b.signMethod = s
+	}
+}
+
 func NewBuilder(root ed25519.PrivateKey, opts ...builderOption) Builder {
 	b := &builder{
 		rng:          rand.Reader,
-		root:         root,
 		symbols:      defaultSymbolTable.Clone(),
 		symbolsStart: defaultSymbolTable.Len(),
 		facts:        new(datalog.FactSet),
+		signMethod:   DefaultSignMethod(root),
 	}
 
 	for _, o := range opts {
@@ -112,7 +127,7 @@ func (b *builder) AddAuthorityCheck(check Check) error {
 }
 
 func (b *builder) Build() (*Biscuit, error) {
-	return New(b.rng, b.root, b.symbols, &Block{
+	return New(b.rng, b.signMethod, b.symbols, &Block{
 		symbols: b.symbols.SplitOff(b.symbolsStart),
 		facts:   b.facts,
 		rules:   b.rules,