ollama/auth/auth.go

62 lines
1.3 KiB
Go
Raw Normal View History

2024-02-05 20:59:52 +00:00
package auth
2023-08-10 18:34:25 +00:00
import (
"bytes"
"context"
2023-08-10 18:34:25 +00:00
"crypto/rand"
"encoding/base64"
"fmt"
"io"
"log/slog"
2023-08-10 18:34:25 +00:00
"os"
2023-09-19 16:36:30 +00:00
"path/filepath"
2023-08-10 18:34:25 +00:00
"golang.org/x/crypto/ssh"
2024-02-05 20:59:52 +00:00
)
2024-02-14 19:29:49 +00:00
const defaultPrivateKey = "id_ed25519"
2023-08-10 18:34:25 +00:00
2024-02-14 19:29:49 +00:00
func NewNonce(r io.Reader, length int) (string, error) {
2023-08-10 18:34:25 +00:00
nonce := make([]byte, length)
2024-02-14 19:29:49 +00:00
if _, err := io.ReadFull(r, nonce); err != nil {
2023-08-10 18:34:25 +00:00
return "", err
}
2023-08-22 01:38:31 +00:00
2024-02-14 19:29:49 +00:00
return base64.RawURLEncoding.EncodeToString(nonce), nil
2023-08-10 18:34:25 +00:00
}
2024-02-14 19:29:49 +00:00
func Sign(ctx context.Context, bts []byte) (string, error) {
2023-08-10 18:34:25 +00:00
home, err := os.UserHomeDir()
if err != nil {
2023-10-20 23:52:48 +00:00
return "", err
2023-08-10 18:34:25 +00:00
}
2024-02-07 19:00:06 +00:00
2024-02-14 19:29:49 +00:00
keyPath := filepath.Join(home, ".ollama", defaultPrivateKey)
2023-08-10 18:34:25 +00:00
2024-02-14 19:29:49 +00:00
privateKeyFile, err := os.ReadFile(keyPath)
2023-08-10 18:34:25 +00:00
if err != nil {
2024-02-14 19:29:49 +00:00
slog.Info(fmt.Sprintf("Failed to load private key: %v", err))
2023-08-10 18:34:25 +00:00
return "", err
}
2024-02-14 19:29:49 +00:00
privateKey, err := ssh.ParsePrivateKey(privateKeyFile)
2023-08-10 18:34:25 +00:00
if err != nil {
return "", err
}
// get the pubkey, but remove the type
2024-02-14 19:29:49 +00:00
publicKey := ssh.MarshalAuthorizedKey(privateKey.PublicKey())
parts := bytes.Split(publicKey, []byte(" "))
2023-08-10 18:34:25 +00:00
if len(parts) < 2 {
return "", fmt.Errorf("malformed public key")
}
2024-02-14 19:29:49 +00:00
signedData, err := privateKey.Sign(rand.Reader, bts)
2023-08-10 18:34:25 +00:00
if err != nil {
return "", err
}
// signature is <pubkey>:<signature>
2024-02-14 19:29:49 +00:00
return fmt.Sprintf("%s:%s", bytes.TrimSpace(parts[1]), base64.StdEncoding.EncodeToString(signedData.Blob)), nil
2023-08-10 18:34:25 +00:00
}