auth.go 1.85 KB
Newer Older
1
package auth
Patrick Devine's avatar
Patrick Devine committed
2
3
4

import (
	"bytes"
5
	"context"
Patrick Devine's avatar
Patrick Devine committed
6
7
	"crypto/rand"
	"encoding/base64"
Michael Yang's avatar
lint  
Michael Yang committed
8
	"errors"
Patrick Devine's avatar
Patrick Devine committed
9
10
	"fmt"
	"io"
11
	"log/slog"
Patrick Devine's avatar
Patrick Devine committed
12
	"os"
Michael Yang's avatar
Michael Yang committed
13
	"path/filepath"
14
	"strings"
Patrick Devine's avatar
Patrick Devine committed
15
16

	"golang.org/x/crypto/ssh"
17
18
)

Michael Yang's avatar
Michael Yang committed
19
const defaultPrivateKey = "id_ed25519"
Patrick Devine's avatar
Patrick Devine committed
20

21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
func keyPath() (string, error) {
	home, err := os.UserHomeDir()
	if err != nil {
		return "", err
	}

	return filepath.Join(home, ".ollama", defaultPrivateKey), nil
}

func GetPublicKey() (string, error) {
	keyPath, err := keyPath()
	if err != nil {
		return "", err
	}

	privateKeyFile, err := os.ReadFile(keyPath)
	if err != nil {
		slog.Info(fmt.Sprintf("Failed to load private key: %v", err))
		return "", err
	}

	privateKey, err := ssh.ParsePrivateKey(privateKeyFile)
	if err != nil {
		return "", err
	}

	publicKey := ssh.MarshalAuthorizedKey(privateKey.PublicKey())

	return strings.TrimSpace(string(publicKey)), nil
}

Michael Yang's avatar
Michael Yang committed
52
func NewNonce(r io.Reader, length int) (string, error) {
Patrick Devine's avatar
Patrick Devine committed
53
	nonce := make([]byte, length)
Michael Yang's avatar
Michael Yang committed
54
	if _, err := io.ReadFull(r, nonce); err != nil {
Patrick Devine's avatar
Patrick Devine committed
55
56
		return "", err
	}
Michael Yang's avatar
Michael Yang committed
57

Michael Yang's avatar
Michael Yang committed
58
	return base64.RawURLEncoding.EncodeToString(nonce), nil
Patrick Devine's avatar
Patrick Devine committed
59
60
}

Michael Yang's avatar
Michael Yang committed
61
func Sign(ctx context.Context, bts []byte) (string, error) {
62
	keyPath, err := keyPath()
Patrick Devine's avatar
Patrick Devine committed
63
	if err != nil {
Michael Yang's avatar
Michael Yang committed
64
		return "", err
Patrick Devine's avatar
Patrick Devine committed
65
	}
Michael Yang's avatar
Michael Yang committed
66

Michael Yang's avatar
Michael Yang committed
67
	privateKeyFile, err := os.ReadFile(keyPath)
Patrick Devine's avatar
Patrick Devine committed
68
	if err != nil {
Michael Yang's avatar
Michael Yang committed
69
		slog.Info(fmt.Sprintf("Failed to load private key: %v", err))
Patrick Devine's avatar
Patrick Devine committed
70
71
72
		return "", err
	}

Michael Yang's avatar
Michael Yang committed
73
	privateKey, err := ssh.ParsePrivateKey(privateKeyFile)
Patrick Devine's avatar
Patrick Devine committed
74
75
76
77
78
	if err != nil {
		return "", err
	}

	// get the pubkey, but remove the type
Michael Yang's avatar
Michael Yang committed
79
80
	publicKey := ssh.MarshalAuthorizedKey(privateKey.PublicKey())
	parts := bytes.Split(publicKey, []byte(" "))
Patrick Devine's avatar
Patrick Devine committed
81
	if len(parts) < 2 {
Michael Yang's avatar
lint  
Michael Yang committed
82
		return "", errors.New("malformed public key")
Patrick Devine's avatar
Patrick Devine committed
83
84
	}

Michael Yang's avatar
Michael Yang committed
85
	signedData, err := privateKey.Sign(rand.Reader, bts)
Patrick Devine's avatar
Patrick Devine committed
86
87
88
89
90
	if err != nil {
		return "", err
	}

	// signature is <pubkey>:<signature>
Michael Yang's avatar
Michael Yang committed
91
	return fmt.Sprintf("%s:%s", bytes.TrimSpace(parts[1]), base64.StdEncoding.EncodeToString(signedData.Blob)), nil
Patrick Devine's avatar
Patrick Devine committed
92
}