parallel_test.go

Documentation: github.com/miekg/pkcs11

     1  package pkcs11
     2  
     3  // A test of using several pkcs11 sessions in parallel for signing across
     4  // multiple goroutines. Access to the PKCS11 module is thread-safe because of
     5  // the C.CKF_OS_LOCKING_OK param and nil mutex functions that the pkcs11
     6  // package passes to C.Initialize, which indicate that the module should use OS
     7  // locking primitives on its own.
     8  //
     9  // Note that while access to the module is thread-safe, sessions are not thread
    10  // safe, and each session must be protected from simultaneous use by some
    11  // synchronization mechanism. In this case we use a cache of sessions (as
    12  // embodied by the `signer` struct), protected by a condition variable. So long
    13  // as there is an available signer in the cache, it is popped off and used. If
    14  // there are no signers available, the caller blocks until there is one
    15  // available.
    16  //
    17  // Please set the appropriate env variables. See the init function.
    18  import (
    19  	"fmt"
    20  	"log"
    21  	"os"
    22  	"sync"
    23  	"testing"
    24  )
    25  
    26  var (
    27  	module          = "/usr/lib/softhsm/libsofthsm.so"
    28  	tokenLabel      = "softhsm token"
    29  	privateKeyLabel = "my key"
    30  	pin             = "1234"
    31  )
    32  
    33  func init() {
    34  	if x := os.Getenv("SOFTHSM_LIB"); x != "" {
    35  		module = x
    36  	}
    37  	if x := os.Getenv("SOFTHSM_TOKENLABEL"); x != "" {
    38  		tokenLabel = x
    39  	}
    40  	if x := os.Getenv("SOFTHSM_PRIVKEYLABEL"); x != "" {
    41  		privateKeyLabel = x
    42  	}
    43  	if x := os.Getenv("SOFTHSM_PIN"); x != "" {
    44  		pin = x
    45  	}
    46  	wd, _ := os.Getwd()
    47  	os.Setenv("SOFTHSM_CONF", wd+"/softhsm.conf")
    48  	os.Setenv("SOFTHSM2_CONF", wd+"/softhsm2.conf")
    49  }
    50  
    51  func initPKCS11Context(modulePath string) (*Ctx, error) {
    52  	context := New(modulePath)
    53  
    54  	if context == nil {
    55  		return nil, fmt.Errorf("unable to load PKCS#11 module")
    56  	}
    57  
    58  	err := context.Initialize()
    59  	return context, err
    60  }
    61  
    62  func getSlot(p *Ctx, label string) (uint, error) {
    63  	slots, err := p.GetSlotList(true)
    64  	if err != nil {
    65  		return 0, err
    66  	}
    67  	for _, slot := range slots {
    68  		_, err := p.GetSlotInfo(slot)
    69  		if err != nil {
    70  			return 0, err
    71  		}
    72  		tokenInfo, err := p.GetTokenInfo(slot)
    73  		if err != nil {
    74  			return 0, err
    75  		}
    76  		if tokenInfo.Label == label {
    77  			return slot, nil
    78  		}
    79  	}
    80  	return 0, fmt.Errorf("Slot not found: %s", label)
    81  }
    82  
    83  func getPrivateKey(context *Ctx, session SessionHandle, label string) (ObjectHandle, error) {
    84  	var noKey ObjectHandle
    85  	template := []*Attribute{
    86  		NewAttribute(CKA_CLASS, CKO_PRIVATE_KEY),
    87  		NewAttribute(CKA_LABEL, label),
    88  	}
    89  	if err := context.FindObjectsInit(session, template); err != nil {
    90  		return noKey, err
    91  	}
    92  	objs, _, err := context.FindObjects(session, 2)
    93  	if err != nil {
    94  		return noKey, err
    95  	}
    96  	if err = context.FindObjectsFinal(session); err != nil {
    97  		return noKey, err
    98  	}
    99  
   100  	if len(objs) == 0 {
   101  		err = fmt.Errorf("private key not found")
   102  		return noKey, err
   103  	}
   104  	return objs[0], nil
   105  }
   106  
   107  type signer struct {
   108  	context    *Ctx
   109  	session    SessionHandle
   110  	privateKey ObjectHandle
   111  }
   112  
   113  func makeSigner(context *Ctx) (*signer, error) {
   114  	slot, err := getSlot(context, tokenLabel)
   115  	if err != nil {
   116  		return nil, err
   117  	}
   118  	session, err := context.OpenSession(slot, CKF_SERIAL_SESSION)
   119  	if err != nil {
   120  		return nil, err
   121  	}
   122  
   123  	if err = context.Login(session, CKU_USER, pin); err != nil {
   124  		context.CloseSession(session)
   125  		return nil, err
   126  	}
   127  
   128  	privateKey, err := getPrivateKey(context, session, privateKeyLabel)
   129  	if err != nil {
   130  		context.CloseSession(session)
   131  		return nil, err
   132  	}
   133  	return &signer{context, session, privateKey}, nil
   134  }
   135  
   136  func (s *signer) sign(input []byte) ([]byte, error) {
   137  	mechanism := []*Mechanism{NewMechanism(CKM_RSA_PKCS, nil)}
   138  	if err := s.context.SignInit(s.session, mechanism, s.privateKey); err != nil {
   139  		log.Fatalf("SignInit: %s", err)
   140  	}
   141  
   142  	signed, err := s.context.Sign(s.session, input)
   143  	if err != nil {
   144  		log.Fatalf("Sign: %s", err)
   145  	}
   146  	return signed, nil
   147  }
   148  
   149  type cache struct {
   150  	signers []*signer
   151  	// this variable signals the condition that there are signers available to be
   152  	// used.
   153  	cond *sync.Cond
   154  }
   155  
   156  func newCache(signers []*signer) cache {
   157  	var mutex sync.Mutex
   158  	return cache{
   159  		signers: signers,
   160  		cond:    sync.NewCond(&mutex),
   161  	}
   162  }
   163  
   164  func (c *cache) get() *signer {
   165  	c.cond.L.Lock()
   166  	for len(c.signers) == 0 {
   167  		c.cond.Wait()
   168  	}
   169  
   170  	instance := c.signers[len(c.signers)-1]
   171  	c.signers = c.signers[:len(c.signers)-1]
   172  	c.cond.L.Unlock()
   173  	return instance
   174  }
   175  
   176  func (c *cache) put(instance *signer) {
   177  	c.cond.L.Lock()
   178  	c.signers = append(c.signers, instance)
   179  	c.cond.Signal()
   180  	c.cond.L.Unlock()
   181  }
   182  
   183  func (c *cache) sign(input []byte) ([]byte, error) {
   184  	instance := c.get()
   185  	defer c.put(instance)
   186  	return instance.sign(input)
   187  }
   188  
   189  // TODO(miek): disabled for now. Fill out the correct values in hsm.db so we can use it.
   190  func testParallel(t *testing.T) {
   191  	if module == "" || tokenLabel == "" || pin == "" || privateKeyLabel == "" {
   192  		t.Fatal("Must pass all flags: module, tokenLabel, pin, and privateKeyLabel")
   193  		return
   194  	}
   195  
   196  	context, err := initPKCS11Context(module)
   197  	if err != nil {
   198  		t.Fatal(err)
   199  	}
   200  
   201  	defer func() {
   202  		context.Finalize()
   203  		context.Destroy()
   204  	}()
   205  
   206  	const nSigners = 100
   207  	const nSignatures = 1000
   208  	signers := make([]*signer, nSigners)
   209  	for i := 0; i < nSigners; i++ {
   210  		signers[i], err = makeSigner(context)
   211  		if err != nil {
   212  			t.Fatalf("Problem making signer: %s", err)
   213  		}
   214  	}
   215  	pool := newCache(signers)
   216  
   217  	output := make(chan []byte, nSignatures)
   218  	for i := 0; i < nSignatures; i++ {
   219  		go func() {
   220  			result, err := pool.sign([]byte("hi"))
   221  			if err != nil {
   222  				t.Fatal(err)
   223  			}
   224  			output <- result
   225  		}()
   226  	}
   227  
   228  	for i := 0; i < nSignatures; i++ {
   229  		// Consume the output of the signers, but do nothing with it.
   230  		<-output
   231  	}
   232  
   233  	for i := 0; i < nSigners; i++ {
   234  		// Note: It is not necessary to call context.Logout. Closing the last
   235  		// session will automatically log out, per PKCS#11 API.
   236  		context.CloseSession(signers[i].session)
   237  	}
   238  }
   239
View as plain text