From d0d5f80a55f07d8ada34777dbfbe5525f9dd6969 Mon Sep 17 00:00:00 2001
From: eyedeekay <hankhill19580@gmail.com>
Date: Mon, 26 May 2025 22:06:42 -0400
Subject: [PATCH] Improve private key parsing

---
 .gitignore      |  1 +
 I2PAddr_test.go | 12 ++++++++++++
 I2PKeys.go      | 28 +++++++++++++++++++++++-----
 NewI2PKeys.go   | 10 ++++++++++
 4 files changed, 46 insertions(+), 5 deletions(-)

diff --git a/.gitignore b/.gitignore
index 6bfe6b1..0af923a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +1,2 @@
 log
+i2p-backup
\ No newline at end of file
diff --git a/I2PAddr_test.go b/I2PAddr_test.go
index e0a3d14..9520248 100644
--- a/I2PAddr_test.go
+++ b/I2PAddr_test.go
@@ -211,6 +211,18 @@ func Test_KeyGenerationAndHandling(t *testing.T) {
 		if buf.String() != expected {
 			t.Errorf("StoreKeysIncompat wrote incorrect data. Got '%s', want '%s'", buf.String(), expected)
 		}
+		// store the buffer content to a permanent local file in this directory
+		err = ioutil.WriteFile("test_keys.txt", buf.Bytes(), 0644)
+		if err != nil {
+			t.Fatalf("Failed to write buffer content to file: '%v'", err)
+		}
+		content, err := ioutil.ReadFile("test_keys.txt")
+		if err != nil {
+			t.Fatalf("Failed to read test_keys.txt: '%v'", err)
+		}
+		if string(content) != expected {
+			t.Errorf("StoreKeysIncompat wrote incorrect data to file. Got '%s', want '%s'", string(content), expected)
+		}
 	})
 
 	t.Run("StoreKeys", func(t *testing.T) {
diff --git a/I2PKeys.go b/I2PKeys.go
index 98623b2..155c330 100644
--- a/I2PKeys.go
+++ b/I2PKeys.go
@@ -164,16 +164,34 @@ func (k I2PKeys) Public() crypto.PublicKey {
 	return k.Address
 }
 
+// Private returns the private key as a byte slice.
 func (k I2PKeys) Private() []byte {
 	log.Debug("Extracting private key")
-	src := strings.Split(k.String(), k.Addr().String())[0]
-	var dest []byte
-	_, err := i2pB64enc.Decode(dest, []byte(src))
+
+	// The private key is everything after the public key in the combined string
+	fullKeys := k.String()
+	publicKey := k.Addr().String()
+
+	// Find where the public key ends in the full string
+	if !strings.HasPrefix(fullKeys, publicKey) {
+		log.Error("Invalid key format: public key not found at start of combined keys")
+		return nil
+	}
+
+	// Extract the private key portion (everything after the public key)
+	privateKeyB64 := fullKeys[len(publicKey):]
+
+	// Pre-allocate destination slice with appropriate capacity
+	dest := make([]byte, i2pB64enc.DecodedLen(len(privateKeyB64)))
+
+	n, err := i2pB64enc.Decode(dest, []byte(privateKeyB64))
 	if err != nil {
 		log.WithError(err).Error("Error decoding private key")
-		panic(err)
+		return nil // Return nil instead of panicking
 	}
-	return dest
+
+	// Return only the portion that was actually decoded
+	return dest[:n]
 }
 
 // Returns the keys (both public and private), in I2Ps base64 format. Use this
diff --git a/NewI2PKeys.go b/NewI2PKeys.go
index 0fb4c58..0e9fcb5 100644
--- a/NewI2PKeys.go
+++ b/NewI2PKeys.go
@@ -116,6 +116,16 @@ func (c *samClient) generateKeys(ctx context.Context, conn net.Conn, keyType str
 	if err != nil {
 		return nil, err
 	}
+	log.Println("Generated keys:", pub, priv)
+	if len(pub) == 0 || len(priv) == 0 {
+		return nil, fmt.Errorf("invalid key response: %s", response)
+	}
+	if len(pub) > maxResponseSize || len(priv) > maxResponseSize {
+		return nil, fmt.Errorf("key response too large: %s", response)
+	}
+	if len(pub) < 128 || len(priv) < 128 {
+		return nil, fmt.Errorf("key response too small: %s", response)
+	}
 
 	return &I2PKeys{
 		Address: I2PAddr(pub),