From 358c5ab877144efa20d3e889826d31d45ea3bec2 Mon Sep 17 00:00:00 2001
From: Ethan Holmes <sherclockholmes@gmail.com>
Date: Sat, 26 Oct 2019 20:43:17 -0700
Subject: [PATCH] Fix private key base64 decode (#30)

* Decode the vapid private key in URL or Raw base64 encoding

* Decode and encode VAPID public key
---
 README.md       |  3 ++-
 example/main.go |  3 ++-
 go.mod          |  2 ++
 vapid.go        | 23 ++++++++++++++++++++---
 webpush.go      |  1 +
 5 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 202f8b0..c313fc6 100644
--- a/README.md
+++ b/README.md
@@ -28,7 +28,7 @@ func main() {
 	json.Unmarshal([]byte("<YOUR_SUBSCRIPTION>"), s)
 
 	// Send Notification
-	_, err := webpush.SendNotification([]byte("Test"), s, &webpush.Options{
+	resp, err := webpush.SendNotification([]byte("Test"), s, &webpush.Options{
 		Subscriber:      "example@example.com",
 		VAPIDPublicKey:  "<YOUR_VAPID_PUBLIC_KEY>",
 		VAPIDPrivateKey: "<YOUR_VAPID_PRIVATE_KEY>",
@@ -37,6 +37,7 @@ func main() {
 	if err != nil {
 		// TODO: Handle error
 	}
+	defer resp.Body.Close()
 }
 ```
 
diff --git a/example/main.go b/example/main.go
index 7d276b1..3445cdb 100644
--- a/example/main.go
+++ b/example/main.go
@@ -18,7 +18,7 @@ func main() {
 	json.Unmarshal([]byte(subscription), s)
 
 	// Send Notification
-	_, err := webpush.SendNotification([]byte("Test"), s, &webpush.Options{
+	resp, err := webpush.SendNotification([]byte("Test"), s, &webpush.Options{
 		Subscriber:      "example@example.com", // Do not include "mailto:"
 		VAPIDPublicKey:  vapidPublicKey,
 		VAPIDPrivateKey: vapidPrivateKey,
@@ -27,4 +27,5 @@ func main() {
 	if err != nil {
 		// TODO: Handle error
 	}
+	defer resp.Body.Close()
 }
diff --git a/go.mod b/go.mod
index 2aa2120..1f2a28f 100644
--- a/go.mod
+++ b/go.mod
@@ -4,3 +4,5 @@ require (
 	github.com/dgrijalva/jwt-go v3.2.0+incompatible
 	golang.org/x/crypto v0.0.0-20190131182504-b8fe1690c613
 )
+
+go 1.13
diff --git a/vapid.go b/vapid.go
index 20dffaa..7aa345f 100644
--- a/vapid.go
+++ b/vapid.go
@@ -77,8 +77,8 @@ func getVAPIDAuthorizationHeader(
 		"sub": fmt.Sprintf("mailto:%s", subscriber),
 	})
 
-	// ECDSA
-	decodedVapidPrivateKey, err := base64.RawURLEncoding.DecodeString(vapidPrivateKey)
+	// Decode the VAPID private key
+	decodedVapidPrivateKey, err := decodeVapidKey(vapidPrivateKey)
 	if err != nil {
 		return "", err
 	}
@@ -91,9 +91,26 @@ func getVAPIDAuthorizationHeader(
 		return "", err
 	}
 
+	// Decode the VAPID public key
+	pubKey, err := decodeVapidKey(vapidPublicKey)
+	if err != nil {
+		return "", err
+	}
+
 	return fmt.Sprintf(
 		"vapid t=%s, k=%s",
 		jwtString,
-		vapidPublicKey,
+		base64.RawURLEncoding.EncodeToString(pubKey),
 	), nil
 }
+
+// Need to decode the vapid private key in multiple base64 formats
+// Solution from: https://github.com/SherClockHolmes/webpush-go/issues/29
+func decodeVapidKey(key string) ([]byte, error) {
+	bytes, err := base64.URLEncoding.DecodeString(key)
+	if err == nil {
+		return bytes, nil
+	}
+
+	return base64.RawURLEncoding.DecodeString(key)
+}
diff --git a/webpush.go b/webpush.go
index 6b41aa0..a31a94d 100644
--- a/webpush.go
+++ b/webpush.go
@@ -229,6 +229,7 @@ func decodeSubscriptionKey(key string) ([]byte, error) {
 	if err == nil {
 		return bytes, nil
 	}
+
 	return base64.URLEncoding.DecodeString(buf.String())
 }