Updated jwt.

mobingilabs · Nov 24, 2017 · 6028274 · 6028274
1 parent 843c4f1
commit 6028274
Show file tree

Hide file tree

Showing 2 changed files with 96 additions and 95 deletions.
diff --git a/pkg/jwt/jwt.go b/pkg/jwt/jwt.go
@@ -1,104 +1,16 @@
 package jwt
 
 import (
-	"crypto/rand"
-	"crypto/rsa"
-	"crypto/x509"
-	"encoding/pem"
 	"fmt"
 	"io/ioutil"
 	"os"
 	"time"
 
 	jwt "github.com/dgrijalva/jwt-go"
 	"github.com/mobingilabs/mobingi-sdk-go/pkg/debug"
-	"github.com/mobingilabs/mobingi-sdk-go/pkg/private"
 	"github.com/pkg/errors"
 )
 
-var (
-	rsainit  bool
-	pubcache []byte
-	prvcache []byte
-	pempub   string
-	pemprv   string
-)
-
-func init() {
-	tmpdir := os.TempDir() + "/jwt/rsa/"
-	pempub = tmpdir + "token.pem.pub"
-	pemprv = tmpdir + "token.pem"
-
-	// create dir if necessary
-	if !private.Exists(tmpdir) {
-		err := os.MkdirAll(tmpdir, 0700)
-		if err != nil {
-			debug.Error(err)
-			return
-		}
-	}
-
-	// create public and private pem files
-	if !private.Exists(pempub) || !private.Exists(pemprv) {
-		priv, err := rsa.GenerateKey(rand.Reader, 2048)
-		if err != nil {
-			debug.Error(err)
-			return
-		}
-
-		privder := x509.MarshalPKCS1PrivateKey(priv)
-		pubkey := priv.Public()
-		pubder, err := x509.MarshalPKIXPublicKey(pubkey)
-		if err != nil {
-			debug.Error(err)
-			return
-		}
-
-		pubblock := &pem.Block{Type: "RSA PUBLIC KEY", Bytes: pubder}
-		pemblock := &pem.Block{Type: "RSA PRIVATE KEY", Bytes: privder}
-		pubfile, err := os.OpenFile(pempub, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
-		if err != nil {
-			debug.Error(err)
-			return
-		}
-
-		defer pubfile.Close()
-		prvfile, err := os.OpenFile(pemprv, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
-		if err != nil {
-			debug.Error(err)
-			return
-		}
-
-		defer prvfile.Close()
-		err = pem.Encode(pubfile, pubblock)
-		if err != nil {
-			debug.Error(err)
-			return
-		}
-
-		err = pem.Encode(prvfile, pemblock)
-		if err != nil {
-			debug.Error(err)
-			return
-		}
-	}
-
-	var err error
-	pubcache, err = ioutil.ReadFile(pempub)
-	if err != nil {
-		debug.Error(err)
-		return
-	}
-
-	prvcache, err = ioutil.ReadFile(pemprv)
-	if err != nil {
-		debug.Error(err)
-		return
-	}
-
-	rsainit = true
-}
-
 type WrapperClaims struct {
 	Data map[string]interface{}
 	jwt.StandardClaims
@@ -109,6 +21,7 @@ type jwtctx struct {
 	Prv    []byte
 	PemPub string
 	PemPrv string
+	init   bool
 }
 
 func (j *jwtctx) GenerateToken(data map[string]interface{}) (*jwt.Token, string, error) {
@@ -147,15 +60,32 @@ func (j *jwtctx) ParseToken(token string) (*jwt.Token, error) {
 	})
 }
 
+// NewCtx initializes our jwt context. For now, it is expected that the pem files
+// (private and public) are already in os.TempDir() + "/jwt/rsa/".
 func NewCtx() (*jwtctx, error) {
-	if !rsainit {
-		return nil, errors.New("failed in rsa init")
+	// TODO: transfer this to authd service
+	tmpdir := os.TempDir() + "/jwt/rsa/"
+	pempub := tmpdir + "token.pem.pub"
+	pemprv := tmpdir + "token.pem"
+
+	pubcache, err := ioutil.ReadFile(pempub)
+	if err != nil {
+		debug.Error(err)
+		return nil, errors.Wrap(err, "pub readfile failed")
+	}
+
+	prvcache, err := ioutil.ReadFile(pemprv)
+	if err != nil {
+		debug.Error(err)
+		return nil, errors.Wrap(err, "prv readfile failed")
+	}
+
+	ctx := jwtctx{
+		PemPub: pempub,
+		PemPrv: pemprv,
+		Pub:    pubcache,
+		Prv:    prvcache,
 	}
 
-	var ctx jwtctx
-	ctx.PemPub = pempub
-	ctx.PemPrv = pemprv
-	ctx.Pub = pubcache
-	ctx.Prv = prvcache
 	return &ctx, nil
 }
diff --git a/pkg/jwt/jwt_test.go b/pkg/jwt/jwt_test.go
@@ -1,12 +1,81 @@
 package jwt
 
 import (
+	"crypto/rand"
+	"crypto/rsa"
+	"crypto/x509"
+	"encoding/pem"
 	"fmt"
 	"log"
+	"os"
 	"testing"
+
+	"github.com/mobingilabs/mobingi-sdk-go/pkg/debug"
+	"github.com/mobingilabs/mobingi-sdk-go/pkg/private"
 )
 
+func initpem() {
+	tmpdir := os.TempDir() + "/jwt/rsa/"
+	pempub := tmpdir + "token.pem.pub"
+	pemprv := tmpdir + "token.pem"
+
+	// create dir if necessary
+	if !private.Exists(tmpdir) {
+		err := os.MkdirAll(tmpdir, 0700)
+		if err != nil {
+			debug.Error(err)
+			return
+		}
+	}
+
+	// create public and private pem files
+	if !private.Exists(pempub) || !private.Exists(pemprv) {
+		priv, err := rsa.GenerateKey(rand.Reader, 2048)
+		if err != nil {
+			debug.Error(err)
+			return
+		}
+
+		privder := x509.MarshalPKCS1PrivateKey(priv)
+		pubkey := priv.Public()
+		pubder, err := x509.MarshalPKIXPublicKey(pubkey)
+		if err != nil {
+			debug.Error(err)
+			return
+		}
+
+		pubblock := &pem.Block{Type: "RSA PUBLIC KEY", Bytes: pubder}
+		pemblock := &pem.Block{Type: "RSA PRIVATE KEY", Bytes: privder}
+		pubfile, err := os.OpenFile(pempub, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
+		if err != nil {
+			debug.Error(err)
+			return
+		}
+
+		defer pubfile.Close()
+		prvfile, err := os.OpenFile(pemprv, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
+		if err != nil {
+			debug.Error(err)
+			return
+		}
+
+		defer prvfile.Close()
+		err = pem.Encode(pubfile, pubblock)
+		if err != nil {
+			debug.Error(err)
+			return
+		}
+
+		err = pem.Encode(prvfile, pemblock)
+		if err != nil {
+			debug.Error(err)
+			return
+		}
+	}
+}
+
 func TestNewCtx(t *testing.T) {
+	initpem()
 	ctx, err := NewCtx()
 	if err != nil {
 		t.Fatal(err)
@@ -18,6 +87,7 @@ func TestNewCtx(t *testing.T) {
 }
 
 func TestGenerateToken(t *testing.T) {
+	initpem()
 	ctx, _ := NewCtx()
 	claims := make(map[string]interface{})
 	claims["username"] = "user"
@@ -30,6 +100,7 @@ func TestGenerateToken(t *testing.T) {
 }
 
 func TestParseToken(t *testing.T) {
+	initpem()
 	ctx, _ := NewCtx()
 	claims := make(map[string]interface{})
 	claims["username"] = "user"