From a5c5cad5764488d9792bca2bef484ff329fa2d4b Mon Sep 17 00:00:00 2001
From: "Jeff R. Allen" <jra@nella.org>
Date: Wed, 22 Jul 2015 01:39:29 +0200
Subject: [PATCH] crypto/openpgp: make it possible to set the key length

Fixes golang/go#6693.

Change-Id: I7322e107bd5f7ad07062dcaadeaa3e85a101015a
Reviewed-on: https://go-review.googlesource.com/12473
Reviewed-by: Adam Langley <agl@golang.org>
Run-TryBot: Adam Langley <agl@golang.org>
---
 keys.go          |  9 +++++++--
 packet/config.go |  3 +++
 write_test.go    | 25 +++++++++++++++++++++++++
 3 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/keys.go b/keys.go
index fe12cfa..bfe3260 100644
--- a/keys.go
+++ b/keys.go
@@ -464,15 +464,20 @@ const defaultRSAKeyBits = 2048
 func NewEntity(name, comment, email string, config *packet.Config) (*Entity, error) {
 	currentTime := config.Now()
 
+	bits := defaultRSAKeyBits
+	if config != nil && config.RSABits != 0 {
+		bits = config.RSABits
+	}
+
 	uid := packet.NewUserId(name, comment, email)
 	if uid == nil {
 		return nil, errors.InvalidArgumentError("user id field contained invalid characters")
 	}
-	signingPriv, err := rsa.GenerateKey(config.Random(), defaultRSAKeyBits)
+	signingPriv, err := rsa.GenerateKey(config.Random(), bits)
 	if err != nil {
 		return nil, err
 	}
-	encryptingPriv, err := rsa.GenerateKey(config.Random(), defaultRSAKeyBits)
+	encryptingPriv, err := rsa.GenerateKey(config.Random(), bits)
 	if err != nil {
 		return nil, err
 	}
diff --git a/packet/config.go b/packet/config.go
index d977cde..c76eecc 100644
--- a/packet/config.go
+++ b/packet/config.go
@@ -43,6 +43,9 @@ type Config struct {
 	// use a value that is at least 65536. See RFC 4880 Section
 	// 3.7.1.3.
 	S2KCount int
+	// RSABits is the number of bits in new RSA keys made with NewEntity.
+	// If zero, then 2048 bit keys are created.
+	RSABits int
 }
 
 func (c *Config) Random() io.Reader {
diff --git a/write_test.go b/write_test.go
index 9f8c358..8e9a335 100644
--- a/write_test.go
+++ b/write_test.go
@@ -10,6 +10,8 @@ import (
 	"io/ioutil"
 	"testing"
 	"time"
+
+	"golang.org/x/crypto/openpgp/packet"
 )
 
 func TestSignDetached(t *testing.T) {
@@ -53,11 +55,34 @@ func TestNewEntity(t *testing.T) {
 		return
 	}
 
+	// Check bit-length with no config.
 	e, err := NewEntity("Test User", "test", "test@example.com", nil)
 	if err != nil {
 		t.Errorf("failed to create entity: %s", err)
 		return
 	}
+	bl, err := e.PrimaryKey.BitLength()
+	if err != nil {
+		t.Errorf("failed to find bit length: %s", err)
+	}
+	if int(bl) != defaultRSAKeyBits {
+		t.Errorf("BitLength %v, expected %v", defaultRSAKeyBits)
+	}
+
+	// Check bit-length with a config.
+	cfg := &packet.Config{RSABits: 1024}
+	e, err = NewEntity("Test User", "test", "test@example.com", cfg)
+	if err != nil {
+		t.Errorf("failed to create entity: %s", err)
+		return
+	}
+	bl, err = e.PrimaryKey.BitLength()
+	if err != nil {
+		t.Errorf("failed to find bit length: %s", err)
+	}
+	if int(bl) != cfg.RSABits {
+		t.Errorf("BitLength %v, expected %v", bl, cfg.RSABits)
+	}
 
 	w := bytes.NewBuffer(nil)
 	if err := e.SerializePrivate(w, nil); err != nil {