core: allow to write documents with AES encryption

pdf/core/crypt.go

@@ -37,7 +37,7 @@ type PdfCrypt struct {
 	U                []byte
 	OE               []byte // R=6
 	UE               []byte // R=6
-	P                int
+	P                int    // TODO (v3): uint32
 	Perms            []byte // R=6
 	EncryptMetadata  bool
 	Id0              string
@@ -97,7 +97,7 @@ const (
 // TODO (v3): Unexport.
 type CryptFilters map[string]CryptFilter
 
-// LoadCryptFilters loads crypt filter information from the encryption dictionary (V4 only).
+// LoadCryptFilters loads crypt filter information from the encryption dictionary (V>=4).
 // TODO (v3): Unexport.
 func (crypt *PdfCrypt) LoadCryptFilters(ed *PdfObjectDictionary) error {
 	crypt.CryptFilters = CryptFilters{}
@@ -213,6 +213,31 @@ func (crypt *PdfCrypt) LoadCryptFilters(ed *PdfObjectDictionary) error {
 	return nil
 }
 
+// SaveCryptFilters saves crypt filter information to the encryption dictionary (V>=4).
+// TODO (v3): Unexport.
+func (crypt *PdfCrypt) SaveCryptFilters(ed *PdfObjectDictionary) error {
+	if crypt.V < 4 {
+		return errors.New("can only be used with V>=4")
+	}
+	cf := MakeDict()
+	ed.Set("CF", cf)
+
+	for name, filter := range crypt.CryptFilters {
+		if name == "Identity" {
+			continue
+		}
+		v := MakeDict()
+		cf.Set(PdfObjectName(name), v)
+
+		v.Set("Type", MakeName("CryptFilter"))
+		v.Set("CFM", MakeName(string(filter.Cfm)))
+		v.Set("Length", MakeInteger(int64(filter.Length)))
+	}
+	ed.Set("StrF", MakeName(crypt.StringFilter))
+	ed.Set("StmF", MakeName(crypt.StreamFilter))
+	return nil
+}
+
 // PdfCryptMakeNew makes the document crypt handler based on the encryption dictionary
 // and trailer dictionary. Returns an error on failure to process.
 func PdfCryptMakeNew(parser *PdfParser, ed, trailer *PdfObjectDictionary) (PdfCrypt, error) {
@@ -273,6 +298,7 @@ func PdfCryptMakeNew(parser *PdfParser, ed, trailer *PdfObjectDictionary) (PdfCr
 	if !ok {
 		return crypter, errors.New("Encrypt dictionary missing R")
 	}
+	// TODO(dennwc): according to spec, R should be validated according to V value
 	if *R < 2 || *R > 6 {
 		return crypter, fmt.Errorf("Invalid R (%d)", *R)
 	}
@@ -1584,10 +1610,24 @@ func (crypt *PdfCrypt) Alg7(opass []byte) (bool, error) {
 	return auth, nil
 }
 
-// encryptR6 is the algorithm opposite to alg2a (R>=5).
+// GenerateParams generates encryption parameters for specified passwords.
+// Can be called only for R>=5.
+func (crypt *PdfCrypt) GenerateParams(upass, opass []byte) error {
+	if crypt.R < 5 {
+		// TODO(dennwc): move code for R<5 from PdfWriter.Encrypt
+		return errors.New("can be used only for R>=5")
+	}
+	crypt.EncryptionKey = make([]byte, 32)
+	if _, err := io.ReadFull(rand.Reader, crypt.EncryptionKey); err != nil {
+		return err
+	}
+	return crypt.generateR6(upass, opass)
+}
+
+// generateR6 is the algorithm opposite to alg2a (R>=5).
 // It generates U,O,UE,OE,Perms fields using AESv3 encryption.
 // There is no algorithm number assigned to this function in the spec.
-func (crypt *PdfCrypt) encryptR6(upass, opass []byte) error {
+func (crypt *PdfCrypt) generateR6(upass, opass []byte) error {
 	// all these field will be populated by functions below
 	crypt.U = nil
 	crypt.O = nil

pdf/core/crypt_test.go

@@ -284,7 +284,7 @@ func TestAESv3(t *testing.T) {
 					}
 
 					// generate encryption parameters
-					err := crypt.encryptR6([]byte(c.UserPass), []byte(c.OwnerPass))
+					err := crypt.generateR6([]byte(c.UserPass), []byte(c.OwnerPass))
 					if err != nil {
 						t.Fatal("Failed to encrypt:", err)
 					}

pdf/model/writer.go

@@ -15,14 +15,15 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"math"
 	"os"
+	"strings"
 	"time"
 
 	"github.com/unidoc/unidoc/common"
 	"github.com/unidoc/unidoc/common/license"
 	. "github.com/unidoc/unidoc/pdf/core"
 	"github.com/unidoc/unidoc/pdf/model/fonts"
-	"strings"
 )
 
 var pdfCreator = ""
@@ -342,8 +343,6 @@ func (this *PdfWriter) AddPage(page *PdfPage) error {
 
 	this.addObject(pageObj)
 
-
-
 	// Traverse the page and record all object references.
 	err := this.addObjects(pDict)
 	if err != nil {
@@ -476,8 +475,21 @@ func (this *PdfWriter) updateObjectNumbers() {
 
 type EncryptOptions struct {
 	Permissions AccessPermissions
+	Algorithm   EncryptionAlgorithm
 }
 
+// EncryptionAlgorithm is used in EncryptOptions to change the default algorithm used to encrypt the document.
+type EncryptionAlgorithm int
+
+const (
+	// RC4_128bit uses RC4 encryption (128 bit)
+	RC4_128bit = EncryptionAlgorithm(iota)
+	// AES_128bit uses AES encryption (128 bit, PDF 1.6)
+	AES_128bit
+	// AES_256bit uses AES encryption (256 bit, PDF 2.0)
+	AES_256bit
+)
+
 // Encrypt the output file with a specified user/owner password.
 func (this *PdfWriter) Encrypt(userPass, ownerPass []byte, options *EncryptOptions) error {
 	crypter := PdfCrypt{}
@@ -486,18 +498,59 @@ func (this *PdfWriter) Encrypt(userPass, ownerPass []byte, options *EncryptOptio
 	crypter.EncryptedObjects = map[PdfObject]bool{}
 
 	crypter.CryptFilters = CryptFilters{}
-	crypter.CryptFilters["Default"] = CryptFilter{Cfm: "V2", Length: 128}
 
-	// Set
+	algo := RC4_128bit
-	crypter.P = -1
+	if options != nil {
+		algo = options.Algorithm
+	}
+
+	var cf CryptFilter
+	switch algo {
+	case RC4_128bit:
 		crypter.V = 2
 		crypter.R = 3
 		crypter.Length = 128
+		cf = CryptFilter{Cfm: CryptFilterV2, Length: 16}
+	case AES_128bit:
+		this.SetVersion(1, 5)
+		crypter.V = 4
+		crypter.R = 4
+		crypter.Length = 128
+		cf = CryptFilter{Cfm: CryptFilterAESV2, Length: 16}
+	case AES_256bit:
+		this.SetVersion(2, 0)
+		crypter.V = 5
+		crypter.R = 6 // TODO(dennwc): a way to set R=5?
+		crypter.Length = 256
+		cf = CryptFilter{Cfm: CryptFilterAESV3, Length: 32}
+	default:
+		return fmt.Errorf("unsupported algorithm: %v", options.Algorithm)
+	}
+	const (
+		defaultFilter = "Default"
+	)
+	crypter.CryptFilters[defaultFilter] = cf
+	if crypter.V >= 4 {
+		crypter.StreamFilter = defaultFilter
+		crypter.StringFilter = defaultFilter
+	}
+
+	// Set
+	crypter.P = math.MaxUint32
 	crypter.EncryptMetadata = true
 	if options != nil {
 		crypter.P = int(options.Permissions.GetP())
 	}
 
+	// Generate the encryption dictionary.
+	ed := MakeDict()
+	ed.Set("Filter", MakeName("Standard"))
+	ed.Set("P", MakeInteger(int64(crypter.P)))
+	ed.Set("V", MakeInteger(int64(crypter.V)))
+	ed.Set("R", MakeInteger(int64(crypter.R)))
+	ed.Set("Length", MakeInteger(int64(crypter.Length)))
+	this.encryptDict = ed
+
 	// Prepare the ID object for the trailer.
 	hashcode := md5.Sum([]byte(time.Now().Format(time.RFC850)))
 	id0 := PdfObjectString(hashcode[:])
@@ -510,6 +563,8 @@ func (this *PdfWriter) Encrypt(userPass, ownerPass []byte, options *EncryptOptio
 	this.ids = &PdfObjectArray{&id0, &id1}
 	common.Log.Trace("Gen Id 0: % x", id0)
 
+	// Generate encryption parameters
+	if crypter.R < 5 {
 		crypter.Id0 = string(id0)
 
 		// Make the O and U objects.
@@ -529,19 +584,30 @@ func (this *PdfWriter) Encrypt(userPass, ownerPass []byte, options *EncryptOptio
 		crypter.U = []byte(U)
 		crypter.EncryptionKey = key
 
-	// Generate the encryption dictionary.
+		ed.Set("O", &O)
-	encDict := MakeDict()
+		ed.Set("U", &U)
-	encDict.Set("Filter", MakeName("Standard"))
+	} else { // R >= 5
-	encDict.Set("P", MakeInteger(int64(crypter.P)))
+		err := crypter.GenerateParams(userPass, ownerPass)
-	encDict.Set("V", MakeInteger(int64(crypter.V)))
+		if err != nil {
-	encDict.Set("R", MakeInteger(int64(crypter.R)))
+			return err
-	encDict.Set("Length", MakeInteger(int64(crypter.Length)))
+		}
-	encDict.Set("O", &O)
+		ed.Set("O", MakeString(string(crypter.O)))
-	encDict.Set("U", &U)
+		ed.Set("U", MakeString(string(crypter.U)))
-	this.encryptDict = encDict
+		ed.Set("OE", MakeString(string(crypter.OE)))
+		ed.Set("UE", MakeString(string(crypter.UE)))
+		ed.Set("EncryptMetadata", MakeBool(crypter.EncryptMetadata))
+		if crypter.R > 5 {
+			ed.Set("Perms", MakeString(string(crypter.Perms)))
+		}
+	}
+	if crypter.V >= 4 {
+		if err := crypter.SaveCryptFilters(ed); err != nil {
+			return err
+		}
+	}
 
-	// Make an object to contain it.
+	// Make an object to contain the encryption dictionary.
-	io := MakeIndirectObject(encDict)
+	io := MakeIndirectObject(ed)
 	this.encryptObj = io
 	this.addObject(io)