cmd: automatically open the browser to register

If a key is not found to belong to any account on ollama.com automatically
open the browser to allow the user to connect their key to an account. This
also outputs a code derived from the key that a user can look at to verify
that they are connecting the device they expect.
This commit is contained in:
Bruce MacDonald
2024-12-02 16:41:02 -08:00
parent ae9165d661
commit 375a662775
3 changed files with 46 additions and 9 deletions

View File

@@ -8,6 +8,7 @@ import (
"crypto/ed25519"
"crypto/rand"
"crypto/sha256"
"encoding/base64"
"encoding/pem"
"errors"
"fmt"
@@ -16,6 +17,7 @@ import (
"math"
"net"
"net/http"
"net/url"
"os"
"os/signal"
"path/filepath"
@@ -30,6 +32,7 @@ import (
"github.com/containerd/console"
"github.com/mattn/go-runewidth"
"github.com/olekukonko/tablewriter"
"github.com/pkg/browser"
"github.com/spf13/cobra"
"golang.org/x/crypto/ssh"
"golang.org/x/term"
@@ -516,7 +519,26 @@ func RunHandler(cmd *cobra.Command, args []string) error {
return generate(cmd, opts)
}
func errFromUnknownKey(unknownKeyErr error) error {
func generateFingerprint(key string) string {
hash := sha256.Sum256([]byte(key))
fingerprint := base64.RawURLEncoding.EncodeToString(hash[:6])
var formatted strings.Builder
for i, char := range fingerprint {
if i > 0 && i%2 == 0 {
formatted.WriteRune('-')
}
formatted.WriteRune(char)
}
return formatted.String()
}
// tryConnect handles key validation when a connection fails due to an unknown key.
// It attempts to open the browser for interactive sessions to let users connect their key,
// falling back to command-line instructions for non-interactive sessions.
// Returns nil if browser flow succeeds, or an error with connection instructions otherwise.
func tryConnect(unknownKeyErr error) error {
// find SSH public key in the error message
// TODO (brucemacd): the API should return structured errors so that this message parsing isn't needed
sshKeyPattern := `ssh-\w+ [^\s"]+`
@@ -545,14 +567,23 @@ func errFromUnknownKey(unknownKeyErr error) error {
return unknownKeyErr
}
var msg strings.Builder
msg.WriteString(unknownKeyErr.Error())
msg.WriteString("\n\nYour ollama key is:\n")
msg.WriteString(localPubKey)
msg.WriteString("\nAdd your key at:\n")
msg.WriteString("https://ollama.com/settings/keys")
if term.IsTerminal(int(os.Stdout.Fd())) {
// URL encode the key and device name for the browser URL
encodedKey := base64.RawURLEncoding.EncodeToString([]byte(localPubKey))
d, _ := os.Hostname()
encodedDevice := url.QueryEscape(d)
browserURL := fmt.Sprintf("https://ollama.com/connect?host=%s&key=%s", encodedDevice, encodedKey)
return errors.New(msg.String())
if err := browser.OpenURL(browserURL); err == nil {
fmt.Printf("\nOpening browser to add your key...\n")
fmt.Printf("\nCheck that this code matches what is shown in your browser:\n")
fmt.Printf("\n %s\n", generateFingerprint(localPubKey))
return nil
}
}
// only return error for non-interactive terminals or if browser opening failed
return fmt.Errorf("%s\nAdd your key at:\nhttps://ollama.com/settings/keys", unknownKeyErr.Error())
}
return unknownKeyErr
@@ -611,13 +642,16 @@ func PushHandler(cmd *cobra.Command, args []string) error {
if spinner != nil {
spinner.Stop()
}
if p != nil {
p.Stop()
}
if strings.Contains(err.Error(), "access denied") {
return errors.New("you are not authorized to push to this namespace, create the model under a namespace you own")
}
if strings.Contains(err.Error(), errtypes.UnknownOllamaKeyErrMsg) && isOllamaHost {
// the user has not added their ollama key to ollama.com
// return an error with a more user-friendly message
return errFromUnknownKey(err)
return tryConnect(err)
}
return err
}