<commit_before>package pkcs11

// Assumption uint is 32 bits on 32 bits platforms and 64 bits on 64 bit platforms

/*
#cgo LDFLAGS: -lltdl
#define CK_PTR *
#ifndef NULL_PTR
#define NULL_PTR 0
#endif
#define CK_DEFINE_FUNCTION(returnType, name) returnType name
#define CK_DECLARE_FUNCTION(returnType, name) returnType name
#define CK_DECLARE_FUNCTION_POINTER(returnType, name) returnType (* name)
#define CK_CALLBACK_FUNCTION(returnType, name) returnType (* name)

#include <stdlib.h>
#include <stdio.h>
#include <ltdl.h>
#include <unistd.h>
#include "pkcs11.h"

struct ctx {
        lt_dlhandle handle;
        CK_FUNCTION_LIST_PTR sym;
};

// New initializes a ctx and fills the symbol table.
struct ctx * New(const char *module) {
        if (lt_dlinit() != 0) {
                return NULL;
        }
        CK_C_GetFunctionList list;
        struct ctx *c=  calloc(1, sizeof(struct ctx));
        c->handle = lt_dlopen(module);
        if (c->handle == NULL) {
                free(c);
                return NULL;
        }
        list = (CK_C_GetFunctionList) lt_dlsym(c->handle, "C_GetFunctionList");
        if (list == NULL) {
                free(c);
                return NULL;
        }
        list(&c->sym);
        return c;
}

// Destroy cleans up a ctx.
void Destroy(struct ctx *c) {
        if (!c) {
                return;
        }
        if (c->handle == NULL) {
                return;
        }
        if (lt_dlclose(c->handle) < 0) {
                return;
        }
        lt_dlexit();
        free(c);
}

CK_RV Initialize(struct ctx* c, CK_VOID_PTR initArgs) {
	return c->sym->C_Initialize(initArgs);
}

CK_RV Finalize(struct ctx* c) {
	return c->sym->C_Finalize(NULL_PTR);
}

CK_RV GetSlotList(struct ctx* c, CK_BBOOL tokenPresent, CK_ULONG_PTR *slotList, CK_ULONG_PTR ulCount) {
	CK_RV e = c->sym->C_GetSlotList(tokenPresent, NULL_PTR, ulCount);
	if (e != CKR_OK) {
		return e;
	}
	*slotList = calloc(1, sizeof(CK_SLOT_ID) * *ulCount);
	e = c->sym->C_GetSlotList(tokenPresent, *slotList, ulCount);
	return e;
}

CK_RV OpenSession(struct ctx* c, CK_ULONG slotID, CK_ULONG flags, CK_SESSION_HANDLE_PTR session) {
	CK_RV e = c->sym->C_OpenSession((CK_SLOT_ID)slotID, (CK_FLAGS)flags, NULL_PTR, NULL_PTR, session);
	return e;
}

CK_RV CloseSession(struct ctx *c, CK_SESSION_HANDLE session) {
	CK_RV e = c->sym->C_CloseSession(session);
	return e;
}

CK_RV Login(struct ctx* c, CK_SESSION_HANDLE session, CK_USER_TYPE userType, char* pin, CK_ULONG pinLen) {
	CK_RV e = c->sym->C_Login(session, userType, (CK_UTF8CHAR_PTR)pin, pinLen);
	return e;
}

CK_RV Logout(struct ctx *c, CK_SESSION_HANDLE session) {
	CK_RV e = c->sym->C_Logout(session);
	return e;
}

CK_RV DigestInit(struct ctx *c, CK_SESSION_HANDLE session, CK_MECHANISM_PTR mechanism) {
	CK_RV e = c->sym->C_DigestInit(session, mechanism);
	return e;
}

CK_RV Digest(struct ctx *c, CK_SESSION_HANDLE session, CK_BYTE_PTR message, CK_ULONG mlen, CK_BYTE_PTR *hash, CK_ULONG_PTR hashlen) {
        CK_RV rv = c->sym->C_Digest(session, message, mlen, NULL, hashlen);
        if (rv != CKR_OK) {
                return rv;
        }
        *hash = malloc(*hashlen * sizeof(CK_BYTE));
	if (*hash == NULL) {
		return CKR_HOST_MEMORY;
	}
        rv = c->sym->C_Digest(session, message, mlen, *hash, hashlen);
	return rv;
}

CK_RV GenerateKeyPair(struct ctx* c, CK_SESSION_HANDLE session, CK_MECHANISM_PTR mechanism,
	CK_ATTRIBUTE_PTR pub, CK_ULONG pubCount, CK_ATTRIBUTE_PTR priv, CK_ULONG privCount,
	CK_OBJECT_HANDLE_PTR pubkey, CK_OBJECT_HANDLE_PTR privkey) {
	CK_RV e = c->sym->C_GenerateKeyPair(session, mechanism, pub, pubCount, priv, privCount,
					pubkey, privkey);
	return e;
}

CK_RV SignInit(struct ctx* c, CK_SESSION_HANDLE session, CK_MECHANISM_PTR mechanism, CK_OBJECT_HANDLE key) {
	CK_RV e = c->sym->C_SignInit(session, mechanism, key);
	return e;
}

CK_RV Sign(struct ctx *c, CK_SESSION_HANDLE session, CK_BYTE_PTR message, CK_ULONG mlen, CK_BYTE_PTR *sig, CK_ULONG_PTR siglen) {
        CK_RV rv = c->sym->C_Sign(session, message, mlen, NULL, siglen);
        if (rv != CKR_OK) {
                return rv;
        }
        *sig = malloc(*siglen * sizeof(CK_BYTE));
	if (*sig == NULL) {
		return CKR_HOST_MEMORY;
	}
        rv = c->sym->C_Sign(session, message, mlen, *sig, siglen);
	return rv;
}

CK_RV EncryptInit(struct ctx* c, CK_SESSION_HANDLE session, CK_MECHANISM_PTR mechanism, CK_OBJECT_HANDLE key) {
	CK_RV e = c->sym->C_EncryptInit(session, mechanism, key);
	return e;
}

CK_RV Encrypt(struct ctx *c, CK_SESSION_HANDLE session, CK_BYTE_PTR message, CK_ULONG mlen, CK_BYTE_PTR *enc, CK_ULONG_PTR enclen) {
        CK_RV rv = c->sym->C_Encrypt(session, message, mlen, NULL, enclen);
        if (rv != CKR_OK) {
                return rv;
        }
        *enc = malloc(*enclen * sizeof(CK_BYTE));
	if (*enc == NULL) {
		return CKR_HOST_MEMORY;
	}
        rv = c->sym->C_Encrypt(session, message, mlen, *enc, enclen);
	return rv;
}

*/
import "C"

import "unsafe"

// Ctx contains the current pkcs11 context.
type Ctx struct {
	ctx         *C.struct_ctx
	initialized bool
	// mutex?
}

// New creates a new context.
func New(module string) *Ctx {
	c := new(Ctx)
	mod := C.CString(module)
	defer C.free(unsafe.Pointer(mod))
	c.ctx = C.New(mod)
	if c.ctx == nil {
		return nil
	}
	return c
}

// Destroy unload the module and frees any remaining memory.
func (c *Ctx) Destroy() {
	if c == nil {
		return
	}
	C.Destroy(c.ctx)
}

func (c *Ctx) Initialize() error {
	args := &C.CK_C_INITIALIZE_ARGS{nil, nil, nil, nil, C.CKF_OS_LOCKING_OK, nil}
	e := C.Initialize(c.ctx, C.CK_VOID_PTR(args))
	if e == C.CKR_OK {
		c.initialized = true // TODO(miek): keep?
	}
	return toError(e)
}

func (c *Ctx) Finalize() error {
	e := C.Finalize(c.ctx)
	return toError(e)
}

func (c *Ctx) GetSlotList(tokenPresent bool) (List, error) {
	var (
		slotList C.CK_ULONG_PTR
		ulCount  C.CK_ULONG
	)
	e := C.GetSlotList(c.ctx, cBBool(tokenPresent), &slotList, &ulCount)
	if toError(e) == nil {
		l := toList(slotList, ulCount)
		return l, nil
	}
	return nil, toError(e)
}

func (c *Ctx) OpenSession(slotID uint, flags uint) (SessionHandle, error) {
	var s C.CK_SESSION_HANDLE
	e := C.OpenSession(c.ctx, C.CK_ULONG(slotID), C.CK_ULONG(flags), C.CK_SESSION_HANDLE_PTR(&s))
	return SessionHandle(s), toError(e)
}

func (c *Ctx) CloseSession(sh SessionHandle) error {
	e := C.CloseSession(c.ctx, C.CK_SESSION_HANDLE(sh))
	return toError(e)
}

func (c *Ctx) Login(sh SessionHandle, userType uint, pin string) error {
	p := C.CString(pin)
	defer C.free(unsafe.Pointer(p))
	e := C.Login(c.ctx, C.CK_SESSION_HANDLE(sh), C.CK_USER_TYPE(userType), p, C.CK_ULONG(len(pin)))
	return toError(e)
}

func (c *Ctx) Logout(sh SessionHandle) error {
	e := C.Logout(c.ctx, C.CK_SESSION_HANDLE(sh))
	return toError(e)
}

func (c *Ctx) GenerateKeyPair(sh SessionHandle, m []*Mechanism, public, private []*Attribute) (ObjectHandle, ObjectHandle, error) {
	var (
		pubkey  C.CK_OBJECT_HANDLE
		privkey C.CK_OBJECT_HANDLE
	)
	pub, pubcount := cAttributeList(public)
	priv, privcount := cAttributeList(private)
	mech, _ := cMechanismList(m)
	e := C.GenerateKeyPair(c.ctx, C.CK_SESSION_HANDLE(sh), mech, pub, pubcount, priv, privcount, C.CK_OBJECT_HANDLE_PTR(&pubkey), C.CK_OBJECT_HANDLE_PTR(&privkey))
	e1 := toError(e)
	if e1 == nil {
		return ObjectHandle(pubkey), ObjectHandle(privkey), nil
	}
	return 0, 0, e1
}

func (c *Ctx) DigestInit(sh SessionHandle, m []*Mechanism) error {
	mech, _ := cMechanismList(m)
	e := C.DigestInit(c.ctx, C.CK_SESSION_HANDLE(sh), mech)
	return toError(e)
}

func (c *Ctx) Digest(sh SessionHandle, message []byte) ([]byte, error) {
	var (
		hash    C.CK_BYTE_PTR
		hashlen C.CK_ULONG
	)
	e := C.Digest(c.ctx, C.CK_SESSION_HANDLE(sh), C.CK_BYTE_PTR(unsafe.Pointer(&message[0])), C.CK_ULONG(len(message)), &hash, &hashlen)
	if toError(e) != nil {
		return nil, toError(e)
	}
	h := C.GoBytes(unsafe.Pointer(hash), C.int(hashlen))
	return h, nil
}

func (c *Ctx) SignInit(sh SessionHandle, m []*Mechanism, o ObjectHandle) error {
	mech, _ := cMechanismList(m) // Only the first is used, but still use a list.
	e := C.SignInit(c.ctx, C.CK_SESSION_HANDLE(sh), mech, C.CK_OBJECT_HANDLE(o))
	return toError(e)
}

func (c *Ctx) Sign(sh SessionHandle, message []byte) ([]byte, error) {
	var (
		sig    C.CK_BYTE_PTR
		siglen C.CK_ULONG
	)
	e := C.Sign(c.ctx, C.CK_SESSION_HANDLE(sh), C.CK_BYTE_PTR(unsafe.Pointer(&message[0])), C.CK_ULONG(len(message)), &sig, &siglen)
	if toError(e) != nil {
		return nil, toError(e)
	}
	s := C.GoBytes(unsafe.Pointer(sig), C.int(siglen))
	return s, nil
}

func (c *Ctx) EncryptInit(sh SessionHandle, m []*Mechanism, o ObjectHandle) error {
	mech, _ := cMechanismList(m)
	e := C.EncryptInit(c.ctx, C.CK_SESSION_HANDLE(sh), mech, C.CK_OBJECT_HANDLE(o))
	return toError(e)
}

func (c *Ctx) Encrypt(sh SessionHandle, message []byte) ([]byte, error) {
	var (
		enc    C.CK_BYTE_PTR
		enclen C.CK_ULONG
	)
	e := C.Encrypt(c.ctx, C.CK_SESSION_HANDLE(sh), C.CK_BYTE_PTR(unsafe.Pointer(&message[0])), C.CK_ULONG(len(message)), &enc, &enclen)
	if toError(e) != nil {
		return nil, toError(e)
	}
	s := C.GoBytes(unsafe.Pointer(enc), C.int(enclen))
	return s, nil
}

<commit_msg>Add CloseAllSessions<commit_after>package pkcs11

// Assumption uint is 32 bits on 32 bits platforms and 64 bits on 64 bit platforms

/*
#cgo LDFLAGS: -lltdl
#define CK_PTR *
#ifndef NULL_PTR
#define NULL_PTR 0
#endif
#define CK_DEFINE_FUNCTION(returnType, name) returnType name
#define CK_DECLARE_FUNCTION(returnType, name) returnType name
#define CK_DECLARE_FUNCTION_POINTER(returnType, name) returnType (* name)
#define CK_CALLBACK_FUNCTION(returnType, name) returnType (* name)

#include <stdlib.h>
#include <stdio.h>
#include <ltdl.h>
#include <unistd.h>
#include "pkcs11.h"

struct ctx {
        lt_dlhandle handle;
        CK_FUNCTION_LIST_PTR sym;
};

// New initializes a ctx and fills the symbol table.
struct ctx * New(const char *module) {
        if (lt_dlinit() != 0) {
                return NULL;
        }
        CK_C_GetFunctionList list;
        struct ctx *c=  calloc(1, sizeof(struct ctx));
        c->handle = lt_dlopen(module);
        if (c->handle == NULL) {
                free(c);
                return NULL;
        }
        list = (CK_C_GetFunctionList) lt_dlsym(c->handle, "C_GetFunctionList");
        if (list == NULL) {
                free(c);
                return NULL;
        }
        list(&c->sym);
        return c;
}

// Destroy cleans up a ctx.
void Destroy(struct ctx *c) {
        if (!c) {
                return;
        }
        if (c->handle == NULL) {
                return;
        }
        if (lt_dlclose(c->handle) < 0) {
                return;
        }
        lt_dlexit();
        free(c);
}

CK_RV Initialize(struct ctx* c, CK_VOID_PTR initArgs) {
	return c->sym->C_Initialize(initArgs);
}

CK_RV Finalize(struct ctx* c) {
	return c->sym->C_Finalize(NULL_PTR);
}

CK_RV GetSlotList(struct ctx* c, CK_BBOOL tokenPresent, CK_ULONG_PTR *slotList, CK_ULONG_PTR ulCount) {
	CK_RV e = c->sym->C_GetSlotList(tokenPresent, NULL_PTR, ulCount);
	if (e != CKR_OK) {
		return e;
	}
	*slotList = calloc(1, sizeof(CK_SLOT_ID) * *ulCount);
	e = c->sym->C_GetSlotList(tokenPresent, *slotList, ulCount);
	return e;
}

CK_RV OpenSession(struct ctx* c, CK_ULONG slotID, CK_ULONG flags, CK_SESSION_HANDLE_PTR session) {
	CK_RV e = c->sym->C_OpenSession((CK_SLOT_ID)slotID, (CK_FLAGS)flags, NULL_PTR, NULL_PTR, session);
	return e;
}

CK_RV CloseSession(struct ctx *c, CK_SESSION_HANDLE session) {
	CK_RV e = c->sym->C_CloseSession(session);
	return e;
}

CK_RV CloseAllSessions(struct ctx *c, CK_ULONG slotID) {
	CK_RV e = c->sym->C_CloseAllSessions(slotID);
	return e;
}

CK_RV Login(struct ctx* c, CK_SESSION_HANDLE session, CK_USER_TYPE userType, char* pin, CK_ULONG pinLen) {
	CK_RV e = c->sym->C_Login(session, userType, (CK_UTF8CHAR_PTR)pin, pinLen);
	return e;
}

CK_RV Logout(struct ctx *c, CK_SESSION_HANDLE session) {
	CK_RV e = c->sym->C_Logout(session);
	return e;
}

CK_RV DigestInit(struct ctx *c, CK_SESSION_HANDLE session, CK_MECHANISM_PTR mechanism) {
	CK_RV e = c->sym->C_DigestInit(session, mechanism);
	return e;
}

CK_RV Digest(struct ctx *c, CK_SESSION_HANDLE session, CK_BYTE_PTR message, CK_ULONG mlen, CK_BYTE_PTR *hash, CK_ULONG_PTR hashlen) {
        CK_RV rv = c->sym->C_Digest(session, message, mlen, NULL, hashlen);
        if (rv != CKR_OK) {
                return rv;
        }
        *hash = malloc(*hashlen * sizeof(CK_BYTE));
	if (*hash == NULL) {
		return CKR_HOST_MEMORY;
	}
        rv = c->sym->C_Digest(session, message, mlen, *hash, hashlen);
	return rv;
}

CK_RV GenerateKeyPair(struct ctx* c, CK_SESSION_HANDLE session, CK_MECHANISM_PTR mechanism,
	CK_ATTRIBUTE_PTR pub, CK_ULONG pubCount, CK_ATTRIBUTE_PTR priv, CK_ULONG privCount,
	CK_OBJECT_HANDLE_PTR pubkey, CK_OBJECT_HANDLE_PTR privkey) {
	CK_RV e = c->sym->C_GenerateKeyPair(session, mechanism, pub, pubCount, priv, privCount,
					pubkey, privkey);
	return e;
}

CK_RV SignInit(struct ctx* c, CK_SESSION_HANDLE session, CK_MECHANISM_PTR mechanism, CK_OBJECT_HANDLE key) {
	CK_RV e = c->sym->C_SignInit(session, mechanism, key);
	return e;
}

CK_RV Sign(struct ctx *c, CK_SESSION_HANDLE session, CK_BYTE_PTR message, CK_ULONG mlen, CK_BYTE_PTR *sig, CK_ULONG_PTR siglen) {
        CK_RV rv = c->sym->C_Sign(session, message, mlen, NULL, siglen);
        if (rv != CKR_OK) {
                return rv;
        }
        *sig = malloc(*siglen * sizeof(CK_BYTE));
	if (*sig == NULL) {
		return CKR_HOST_MEMORY;
	}
        rv = c->sym->C_Sign(session, message, mlen, *sig, siglen);
	return rv;
}

CK_RV EncryptInit(struct ctx* c, CK_SESSION_HANDLE session, CK_MECHANISM_PTR mechanism, CK_OBJECT_HANDLE key) {
	CK_RV e = c->sym->C_EncryptInit(session, mechanism, key);
	return e;
}

CK_RV Encrypt(struct ctx *c, CK_SESSION_HANDLE session, CK_BYTE_PTR message, CK_ULONG mlen, CK_BYTE_PTR *enc, CK_ULONG_PTR enclen) {
        CK_RV rv = c->sym->C_Encrypt(session, message, mlen, NULL, enclen);
        if (rv != CKR_OK) {
                return rv;
        }
        *enc = malloc(*enclen * sizeof(CK_BYTE));
	if (*enc == NULL) {
		return CKR_HOST_MEMORY;
	}
        rv = c->sym->C_Encrypt(session, message, mlen, *enc, enclen);
	return rv;
}

*/
import "C"

import "unsafe"

// Ctx contains the current pkcs11 context.
type Ctx struct {
	ctx         *C.struct_ctx
	initialized bool
	// mutex?
}

// New creates a new context.
func New(module string) *Ctx {
	c := new(Ctx)
	mod := C.CString(module)
	defer C.free(unsafe.Pointer(mod))
	c.ctx = C.New(mod)
	if c.ctx == nil {
		return nil
	}
	return c
}

// Destroy unload the module and frees any remaining memory.
func (c *Ctx) Destroy() {
	if c == nil {
		return
	}
	C.Destroy(c.ctx)
}

func (c *Ctx) Initialize() error {
	args := &C.CK_C_INITIALIZE_ARGS{nil, nil, nil, nil, C.CKF_OS_LOCKING_OK, nil}
	e := C.Initialize(c.ctx, C.CK_VOID_PTR(args))
	if e == C.CKR_OK {
		c.initialized = true // TODO(miek): keep?
	}
	return toError(e)
}

func (c *Ctx) Finalize() error {
	e := C.Finalize(c.ctx)
	return toError(e)
}

func (c *Ctx) GetSlotList(tokenPresent bool) (List, error) {
	var (
		slotList C.CK_ULONG_PTR
		ulCount  C.CK_ULONG
	)
	e := C.GetSlotList(c.ctx, cBBool(tokenPresent), &slotList, &ulCount)
	if toError(e) == nil {
		l := toList(slotList, ulCount)
		return l, nil
	}
	return nil, toError(e)
}

func (c *Ctx) OpenSession(slotID uint, flags uint) (SessionHandle, error) {
	var s C.CK_SESSION_HANDLE
	e := C.OpenSession(c.ctx, C.CK_ULONG(slotID), C.CK_ULONG(flags), C.CK_SESSION_HANDLE_PTR(&s))
	return SessionHandle(s), toError(e)
}

func (c *Ctx) CloseSession(sh SessionHandle) error {
	e := C.CloseSession(c.ctx, C.CK_SESSION_HANDLE(sh))
	return toError(e)
}

func (c *Ctx) CloseAllSessions(slotID uint) error {
	e := C.CloseAllSessions(c.ctx, C.CK_ULONG(slotID))
	return toError(e)
}

func (c *Ctx) Login(sh SessionHandle, userType uint, pin string) error {
	p := C.CString(pin)
	defer C.free(unsafe.Pointer(p))
	e := C.Login(c.ctx, C.CK_SESSION_HANDLE(sh), C.CK_USER_TYPE(userType), p, C.CK_ULONG(len(pin)))
	return toError(e)
}

func (c *Ctx) Logout(sh SessionHandle) error {
	e := C.Logout(c.ctx, C.CK_SESSION_HANDLE(sh))
	return toError(e)
}

func (c *Ctx) GenerateKeyPair(sh SessionHandle, m []*Mechanism, public, private []*Attribute) (ObjectHandle, ObjectHandle, error) {
	var (
		pubkey  C.CK_OBJECT_HANDLE
		privkey C.CK_OBJECT_HANDLE
	)
	pub, pubcount := cAttributeList(public)
	priv, privcount := cAttributeList(private)
	mech, _ := cMechanismList(m)
	e := C.GenerateKeyPair(c.ctx, C.CK_SESSION_HANDLE(sh), mech, pub, pubcount, priv, privcount, C.CK_OBJECT_HANDLE_PTR(&pubkey), C.CK_OBJECT_HANDLE_PTR(&privkey))
	e1 := toError(e)
	if e1 == nil {
		return ObjectHandle(pubkey), ObjectHandle(privkey), nil
	}
	return 0, 0, e1
}

func (c *Ctx) DigestInit(sh SessionHandle, m []*Mechanism) error {
	mech, _ := cMechanismList(m)
	e := C.DigestInit(c.ctx, C.CK_SESSION_HANDLE(sh), mech)
	return toError(e)
}

func (c *Ctx) Digest(sh SessionHandle, message []byte) ([]byte, error) {
	var (
		hash    C.CK_BYTE_PTR
		hashlen C.CK_ULONG
	)
	e := C.Digest(c.ctx, C.CK_SESSION_HANDLE(sh), C.CK_BYTE_PTR(unsafe.Pointer(&message[0])), C.CK_ULONG(len(message)), &hash, &hashlen)
	if toError(e) != nil {
		return nil, toError(e)
	}
	h := C.GoBytes(unsafe.Pointer(hash), C.int(hashlen))
	return h, nil
}

func (c *Ctx) SignInit(sh SessionHandle, m []*Mechanism, o ObjectHandle) error {
	mech, _ := cMechanismList(m) // Only the first is used, but still use a list.
	e := C.SignInit(c.ctx, C.CK_SESSION_HANDLE(sh), mech, C.CK_OBJECT_HANDLE(o))
	return toError(e)
}

func (c *Ctx) Sign(sh SessionHandle, message []byte) ([]byte, error) {
	var (
		sig    C.CK_BYTE_PTR
		siglen C.CK_ULONG
	)
	e := C.Sign(c.ctx, C.CK_SESSION_HANDLE(sh), C.CK_BYTE_PTR(unsafe.Pointer(&message[0])), C.CK_ULONG(len(message)), &sig, &siglen)
	if toError(e) != nil {
		return nil, toError(e)
	}
	s := C.GoBytes(unsafe.Pointer(sig), C.int(siglen))
	return s, nil
}

func (c *Ctx) EncryptInit(sh SessionHandle, m []*Mechanism, o ObjectHandle) error {
	mech, _ := cMechanismList(m)
	e := C.EncryptInit(c.ctx, C.CK_SESSION_HANDLE(sh), mech, C.CK_OBJECT_HANDLE(o))
	return toError(e)
}

func (c *Ctx) Encrypt(sh SessionHandle, message []byte) ([]byte, error) {
	var (
		enc    C.CK_BYTE_PTR
		enclen C.CK_ULONG
	)
	e := C.Encrypt(c.ctx, C.CK_SESSION_HANDLE(sh), C.CK_BYTE_PTR(unsafe.Pointer(&message[0])), C.CK_ULONG(len(message)), &enc, &enclen)
	if toError(e) != nil {
		return nil, toError(e)
	}
	s := C.GoBytes(unsafe.Pointer(enc), C.int(enclen))
	return s, nil
}

<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"net"
	"os"
	"os/signal"
	"os/user"
	"syscall"

	arg "github.com/alexflint/go-arg"
	"github.com/dshearer/jobber/common"
)

var gUser *user.User
var gIpcServer IpcServer
var gJobManager *JobManager

func quit(exitCode int) {
	if gIpcServer != nil {
		gIpcServer.Stop()
	}
	if gJobManager != nil {
		gJobManager.Cancel()
		gJobManager.Wait()
	}
	os.Exit(exitCode)
}

func quitOnSignal() {
	// Set up channel on which to send signal notifications.
	// We must use a buffered channel or risk missing the signal
	// if we're not ready to receive when the signal is sent.
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM, os.Kill)

	// Block until a signal is received.
	<-c
	quit(0)
}

func quitOnJobbermasterDiscon(quickSockPath string) {
	/*
	   Jobbermaster launched us and gave us a path to a Unix socket.
	   When jobbermaster wants us to quit, it will close that socket.
	*/

	common.Logger.Printf(
		"jobbermaster quit socket: %v",
		quickSockPath,
	)

	// open socket
	addr, err := net.ResolveUnixAddr("unix", quickSockPath)
	if err != nil {
		common.ErrLogger.Printf(
			"ResolveUnixAddr failed on %v: %v",
			quickSockPath,
			err,
		)
		quit(1)
		return
	}
	conn, err := net.DialUnix("unix", nil, addr)
	if err != nil {
		common.ErrLogger.Printf(
			"DialUnix failed on %v: %v",
			quickSockPath,
			err,
		)
		quit(1)
		return
	}
	defer conn.Close()

	// read from it -- this will block until jobbermaster close the cxn
	var buf [1]byte
	_, _ = conn.Read(buf[:])

	/*
	   As conn.Read has returned, jobbermaster must have closed the cxn.
	*/
	quit(0)
}

func usage(extraMsg string) {
	if len(extraMsg) > 0 {
		fmt.Printf("%v\n\n", extraMsg)
	}

	fmt.Printf("Usage: %v [flags] JOBFILE_PATH\n\n", os.Args[0])

	fmt.Printf("Flags:\n")
	flag.PrintDefaults()
	fmt.Printf("\n")
}

type argsS struct {
	QuitSocket  *string `arg:"-q" help:"path to quit socket (used by jobbermaster to tell us to quit)"`
	UnixSocket  *string `arg:"-u" help:"path to Unix socket on which to receive commands"`
	TcpPort     *uint   `arg:"-p" help:"TCP port on which to receive commands"`
	TempDir     *string `arg:"-t" help:"Path to dir to use as temp dir"`
	JobfilePath string  `arg:"positional,required"`
	Debug       bool    `arg:"-d" default:"false"`
}

func (argsS) Version() string {
	return common.LongVersionStr()
}

func main() {

	// parse args
	var args argsS
	arg.MustParse(&args)

	// check for errors
	if (args.UnixSocket == nil) == (args.TcpPort == nil) {
		fmt.Fprintf(os.Stderr, "Must specify exactly one of --unixsocket or --tcpport\n")
		quit(1)
	}
	if args.UnixSocket != nil && len(*args.UnixSocket) == 0 {
		fmt.Fprintf(os.Stderr, "Unix socket path cannot be empty\n")
		quit(1)
	}
	if args.QuitSocket != nil && len(*args.QuitSocket) == 0 {
		fmt.Fprintf(os.Stderr, "Quit socket path cannot be empty\n")
		quit(1)
	}
	if args.TempDir != nil && len(*args.TempDir) == 0 {
		fmt.Fprintf(os.Stderr, "Temp dir path cannot be empty\n")
		quit(1)
	}
	if args.TcpPort != nil && *args.TcpPort == 0 {
		fmt.Fprintf(os.Stderr, "TCP port cannot be zero\n")
		quit(1)
	}
	if len(args.JobfilePath) == 0 {
		fmt.Fprintf(os.Stderr, "Jobfile path cannot be empty")
		quit(1)
	}

	// init settings
	err := common.InitSettings(common.InitSettingsParams{
		TempDir: args.TempDir,
	})
	if err != nil {
		common.ErrLogger.Print(err)
		os.Exit(1)
	}

	if args.Debug {
		common.PrintPaths()
		os.Exit(0)
	}

	// get current user
	gUser, err = user.Current()
	if err != nil {
		common.ErrLogger.Printf("Failed to get current user: %v", err)
		quit(1)
	}

	// run job manager
	gJobManager = NewJobManager(args.JobfilePath)
	if err := gJobManager.Launch(); err != nil {
		common.ErrLogger.Printf("Error: %v\n", err)
		quit(1)
	}

	// make IPC server
	if args.UnixSocket != nil {
		gIpcServer = NewUdsIpcServer(*args.UnixSocket, gJobManager.CmdChan)
		common.Logger.Printf("Listening for commands on %v", *args.UnixSocket)
	} else {
		gIpcServer = NewInetIpcServer(*args.TcpPort, gJobManager.CmdChan)
		common.Logger.Printf("Listening for commands on :%v", *args.TcpPort)
	}
	if err := gIpcServer.Launch(); err != nil {
		common.ErrLogger.Printf("Error: %v", err)
		quit(1)
	}

	if args.QuitSocket != nil {
		// listen for jobbermaster to tell us to quit
		go quitOnJobbermasterDiscon(*args.QuitSocket)
	}

	// listen for signals
	go quitOnSignal()

	// wait for job manager
	gJobManager.Wait()

	quit(0)
}
<commit_msg>jobberrunner: Specifying socket for IPC is no longer required<commit_after>package main

import (
	"flag"
	"fmt"
	"net"
	"os"
	"os/signal"
	"os/user"
	"syscall"

	arg "github.com/alexflint/go-arg"
	"github.com/dshearer/jobber/common"
)

var gUser *user.User
var gIpcServer IpcServer
var gJobManager *JobManager

func quit(exitCode int) {
	if gIpcServer != nil {
		gIpcServer.Stop()
	}
	if gJobManager != nil {
		gJobManager.Cancel()
		gJobManager.Wait()
	}
	os.Exit(exitCode)
}

func quitOnSignal() {
	// Set up channel on which to send signal notifications.
	// We must use a buffered channel or risk missing the signal
	// if we're not ready to receive when the signal is sent.
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM, os.Kill)

	// Block until a signal is received.
	<-c
	quit(0)
}

func quitOnJobbermasterDiscon(quickSockPath string) {
	/*
	   Jobbermaster launched us and gave us a path to a Unix socket.
	   When jobbermaster wants us to quit, it will close that socket.
	*/

	common.Logger.Printf(
		"jobbermaster quit socket: %v",
		quickSockPath,
	)

	// open socket
	addr, err := net.ResolveUnixAddr("unix", quickSockPath)
	if err != nil {
		common.ErrLogger.Printf(
			"ResolveUnixAddr failed on %v: %v",
			quickSockPath,
			err,
		)
		quit(1)
		return
	}
	conn, err := net.DialUnix("unix", nil, addr)
	if err != nil {
		common.ErrLogger.Printf(
			"DialUnix failed on %v: %v",
			quickSockPath,
			err,
		)
		quit(1)
		return
	}
	defer conn.Close()

	// read from it -- this will block until jobbermaster close the cxn
	var buf [1]byte
	_, _ = conn.Read(buf[:])

	/*
	   As conn.Read has returned, jobbermaster must have closed the cxn.
	*/
	quit(0)
}

func usage(extraMsg string) {
	if len(extraMsg) > 0 {
		fmt.Printf("%v\n\n", extraMsg)
	}

	fmt.Printf("Usage: %v [flags] JOBFILE_PATH\n\n", os.Args[0])

	fmt.Printf("Flags:\n")
	flag.PrintDefaults()
	fmt.Printf("\n")
}

type argsS struct {
	QuitSocket  *string `arg:"-q" help:"path to quit socket (used by jobbermaster to tell us to quit)"`
	UnixSocket  *string `arg:"-u" help:"path to Unix socket on which to receive commands"`
	TcpPort     *uint   `arg:"-p" help:"TCP port on which to receive commands"`
	TempDir     *string `arg:"-t" help:"Path to dir to use as temp dir"`
	JobfilePath string  `arg:"positional,required"`
	Debug       bool    `arg:"-d" default:"false"`
}

func (argsS) Version() string {
	return common.LongVersionStr()
}

func main() {

	// parse args
	var args argsS
	arg.MustParse(&args)

	// check for errors
	if args.UnixSocket != nil && args.TcpPort != nil {
		fmt.Fprintf(os.Stderr, "Must specify at most one of --unixsocket or --tcpport\n")
		quit(1)
	}
	if args.UnixSocket != nil && len(*args.UnixSocket) == 0 {
		fmt.Fprintf(os.Stderr, "Unix socket path cannot be empty\n")
		quit(1)
	}
	if args.QuitSocket != nil && len(*args.QuitSocket) == 0 {
		fmt.Fprintf(os.Stderr, "Quit socket path cannot be empty\n")
		quit(1)
	}
	if args.TempDir != nil && len(*args.TempDir) == 0 {
		fmt.Fprintf(os.Stderr, "Temp dir path cannot be empty\n")
		quit(1)
	}
	if args.TcpPort != nil && *args.TcpPort == 0 {
		fmt.Fprintf(os.Stderr, "TCP port cannot be zero\n")
		quit(1)
	}
	if len(args.JobfilePath) == 0 {
		fmt.Fprintf(os.Stderr, "Jobfile path cannot be empty")
		quit(1)
	}

	// init settings
	err := common.InitSettings(common.InitSettingsParams{
		TempDir: args.TempDir,
	})
	if err != nil {
		common.ErrLogger.Print(err)
		os.Exit(1)
	}

	if args.Debug {
		common.PrintPaths()
		os.Exit(0)
	}

	// get current user
	gUser, err = user.Current()
	if err != nil {
		common.ErrLogger.Printf("Failed to get current user: %v", err)
		quit(1)
	}

	// run job manager
	gJobManager = NewJobManager(args.JobfilePath)
	if err := gJobManager.Launch(); err != nil {
		common.ErrLogger.Printf("Error: %v\n", err)
		quit(1)
	}

	// make IPC server
	if args.UnixSocket != nil {
		gIpcServer = NewUdsIpcServer(*args.UnixSocket, gJobManager.CmdChan)
		common.Logger.Printf("Listening for commands on %v", *args.UnixSocket)
	} else if args.TcpPort != nil {
		gIpcServer = NewInetIpcServer(*args.TcpPort, gJobManager.CmdChan)
		common.Logger.Printf("Listening for commands on :%v", *args.TcpPort)
	}
	if gIpcServer != nil {
		if err := gIpcServer.Launch(); err != nil {
			common.ErrLogger.Printf("Error: %v", err)
			quit(1)
		}
	}

	if args.QuitSocket != nil {
		// listen for jobbermaster to tell us to quit
		go quitOnJobbermasterDiscon(*args.QuitSocket)
	}

	// listen for signals
	go quitOnSignal()

	// wait for job manager
	gJobManager.Wait()

	quit(0)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sts"
)

type (
	// Config holds input parameters for the plugin
	Config struct {
		Actions          []string
		Vars             map[string]string
		Secrets          map[string]string
		InitOptions      InitOptions
		FmtOptions       FmtOptions
		Cacert           string
		Sensitive        bool
		RoleARN          string
		RootDir          string
		Parallelism      int
		Targets          []string
		VarFiles         []string
		TerraformDataDir string
		DisableRefresh   bool
	}

	// Netrc is credentials for cloning
	Netrc struct {
		Machine  string
		Login    string
		Password string
	}

	// InitOptions include options for the Terraform's init command
	InitOptions struct {
		BackendConfig []string `json:"backend-config"`
		Lock          *bool    `json:"lock"`
		LockTimeout   string   `json:"lock-timeout"`
	}

	// FmtOptions fmt options for the Terraform's fmt command
	FmtOptions struct {
		List  *bool `json:"list"`
		Write *bool `json:"write"`
		Diff  *bool `json:"diff"`
		Check *bool `json:"check"`
	}

	// Plugin represents the plugin instance to be executed
	Plugin struct {
		Config    Config
		Netrc     Netrc
		Terraform Terraform
	}
)

// Exec executes the plugin
func (p Plugin) Exec() error {
	// Install specified version of terraform
	if p.Terraform.Version != "" {
		err := installTerraform(p.Terraform.Version)

		if err != nil {
			return err
		}
	}

	if p.Config.RoleARN != "" && !credsSet() {
		assumeRole(p.Config.RoleARN)
	}

	// writing the .netrc file with Github credentials in it.
	err := writeNetrc(p.Netrc.Machine, p.Netrc.Login, p.Netrc.Password)
	if err != nil {
		return err
	}

	terraformDataDir := ".terraform"
	if p.Config.TerraformDataDir != "" {
		terraformDataDir = p.Config.TerraformDataDir
		os.Setenv("TF_DATA_DIR", p.Config.TerraformDataDir)
	}

	var commands []*exec.Cmd

	commands = append(commands, exec.Command("terraform", "version"))

	CopyTfEnv()

	if p.Config.Cacert != "" {
		commands = append(commands, installCaCert(p.Config.Cacert))
	}

	commands = append(commands, deleteCache(terraformDataDir))
	commands = append(commands, initCommand(p.Config.InitOptions))
	commands = append(commands, getModules())

	// Add commands listed from Actions
	for _, action := range p.Config.Actions {
		switch action {
		case "fmt":
			commands = append(commands, tfFmt(p.Config))
		case "validate":
			commands = append(commands, tfValidate())
		case "plan":
			commands = append(commands, tfPlan(p.Config, false))
		case "plan-destroy":
			commands = append(commands, tfPlan(p.Config, true))
		case "apply":
			commands = append(commands, tfApply(p.Config))
		case "destroy":
			commands = append(commands, tfDestroy(p.Config))
		default:
			return fmt.Errorf("valid actions are: fmt, validate, plan, apply, plan-destroy, destroy.  You provided %s", action)
		}
	}

	commands = append(commands, deleteCache(terraformDataDir))

	for _, c := range commands {
		if c.Dir == "" {
			wd, err := os.Getwd()
			if err == nil {
				c.Dir = wd
			}
		}
		if p.Config.RootDir != "" {
			c.Dir = c.Dir + "/" + p.Config.RootDir
		}
		c.Stdout = os.Stdout
		c.Stderr = os.Stderr
		if !p.Config.Sensitive {
			trace(c)
		}

		err := c.Run()
		if err != nil {
			logrus.WithFields(logrus.Fields{
				"error": err,
			}).Fatal("Failed to execute a command")
		}
		logrus.Debug("Command completed successfully")
	}

	return nil
}

// CopyTfEnv creates copies of TF_VAR_ to lowercase
func CopyTfEnv() {
	tfVar := regexp.MustCompile(`^TF_VAR_.*$`)
	for _, e := range os.Environ() {
		pair := strings.SplitN(e, "=", 2)
		if tfVar.MatchString(pair[0]) {
			name := strings.Split(pair[0], "TF_VAR_")
			os.Setenv(fmt.Sprintf("TF_VAR_%s", strings.ToLower(name[1])), pair[1])
		}
	}
}

func credsSet() bool {
	awsTokens := []string{"AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN"}
	for _, token := range awsTokens {
		if os.Getenv(token) == "" {
			return false
		}
	}
	return true
}

func assumeRole(roleArn string) {
	client := sts.New(session.New())
	duration := time.Hour * 1
	stsProvider := &stscreds.AssumeRoleProvider{
		Client:          client,
		Duration:        duration,
		RoleARN:         roleArn,
		RoleSessionName: "drone",
	}

	value, err := credentials.NewCredentials(stsProvider).Get()
	if err != nil {
		logrus.WithFields(logrus.Fields{
			"error": err,
		}).Fatal("Error assuming role!")
	}
	os.Setenv("AWS_ACCESS_KEY_ID", value.AccessKeyID)
	os.Setenv("AWS_SECRET_ACCESS_KEY", value.SecretAccessKey)
	os.Setenv("AWS_SESSION_TOKEN", value.SessionToken)
}

func deleteCache(terraformDataDir string) *exec.Cmd {
	return exec.Command(
		"rm",
		"-rf",
		terraformDataDir,
	)
}

func getModules() *exec.Cmd {
	return exec.Command(
		"terraform",
		"get",
	)
}

func initCommand(config InitOptions) *exec.Cmd {
	args := []string{
		"init",
	}

	for _, v := range config.BackendConfig {
		args = append(args, fmt.Sprintf("-backend-config=%s", v))
	}

	// True is default in TF
	if config.Lock != nil {
		args = append(args, fmt.Sprintf("-lock=%t", *config.Lock))
	}

	// "0s" is default in TF
	if config.LockTimeout != "" {
		args = append(args, fmt.Sprintf("-lock-timeout=%s", config.LockTimeout))
	}

	// Fail Terraform execution on prompt
	args = append(args, "-input=false")

	return exec.Command(
		"terraform",
		args...,
	)
}

func installCaCert(cacert string) *exec.Cmd {
	ioutil.WriteFile("/usr/local/share/ca-certificates/ca_cert.crt", []byte(cacert), 0644)
	return exec.Command(
		"update-ca-certificates",
	)
}

func trace(cmd *exec.Cmd) {
	fmt.Println("$", strings.Join(cmd.Args, " "))
}

func tfApply(config Config) *exec.Cmd {
	args := []string{
		"apply",
	}
	for _, v := range config.Targets {
		args = append(args, "--target", fmt.Sprintf("%s", v))
	}
	if config.Parallelism > 0 {
		args = append(args, fmt.Sprintf("-parallelism=%d", config.Parallelism))
	}
	if config.InitOptions.Lock != nil {
		args = append(args, fmt.Sprintf("-lock=%t", *config.InitOptions.Lock))
	}
	if config.InitOptions.LockTimeout != "" {
		args = append(args, fmt.Sprintf("-lock-timeout=%s", config.InitOptions.LockTimeout))
	}
	if config.DisableRefresh {
		args = append(args, "-refresh=false")
	}
	args = append(args, getTfoutPath())

	return exec.Command(
		"terraform",
		args...,
	)
}

func tfDestroy(config Config) *exec.Cmd {
	args := []string{
		"destroy",
	}
	for _, v := range config.Targets {
		args = append(args, fmt.Sprintf("-target=%s", v))
	}
	args = append(args, varFiles(config.VarFiles)...)
	args = append(args, vars(config.Vars)...)
	if config.Parallelism > 0 {
		args = append(args, fmt.Sprintf("-parallelism=%d", config.Parallelism))
	}
	if config.InitOptions.Lock != nil {
		args = append(args, fmt.Sprintf("-lock=%t", *config.InitOptions.Lock))
	}
	if config.InitOptions.LockTimeout != "" {
		args = append(args, fmt.Sprintf("-lock-timeout=%s", config.InitOptions.LockTimeout))
	}
	args = append(args, "-force")
	return exec.Command(
		"terraform",
		args...,
	)
}

func tfPlan(config Config, destroy bool) *exec.Cmd {
	args := []string{
		"plan",
	}

	if destroy {
		args = append(args, "-destroy")
	} else {
		args = append(args, fmt.Sprintf("-out=%s", getTfoutPath()))
	}

	for _, v := range config.Targets {
		args = append(args, "--target", fmt.Sprintf("%s", v))
	}
	args = append(args, varFiles(config.VarFiles)...)
	args = append(args, vars(config.Vars)...)
	if config.Parallelism > 0 {
		args = append(args, fmt.Sprintf("-parallelism=%d", config.Parallelism))
	}
	if config.InitOptions.Lock != nil {
		args = append(args, fmt.Sprintf("-lock=%t", *config.InitOptions.Lock))
	}
	if config.InitOptions.LockTimeout != "" {
		args = append(args, fmt.Sprintf("-lock-timeout=%s", config.InitOptions.LockTimeout))
	}
	if config.DisableRefresh {
		args = append(args, "-refresh=false")
	}
	return exec.Command(
		"terraform",
		args...,
	)
}

func tfValidate() *exec.Cmd {
	return exec.Command(
		"terraform",
		"validate",
	)
}

func tfFmt(config Config) *exec.Cmd {
	args := []string{
		"fmt",
	}
	if config.FmtOptions.List != nil {
		args = append(args, fmt.Sprintf("-list=%t", *config.FmtOptions.List))
	}
	if config.FmtOptions.Write != nil {
		args = append(args, fmt.Sprintf("-write=%t", *config.FmtOptions.Write))
	}
	if config.FmtOptions.Diff != nil {
		args = append(args, fmt.Sprintf("-diff=%t", *config.FmtOptions.Diff))
	}
	if config.FmtOptions.Check != nil {
		args = append(args, fmt.Sprintf("-check=%t", *config.FmtOptions.Check))
	}
	return exec.Command(
		"terraform",
		args...,
	)
}

func getTfoutPath() string {
	terraformDataDir := os.Getenv("TF_DATA_DIR")
	if terraformDataDir == ".terraform" || terraformDataDir == "" {
		return "plan.tfout"
	}

	return fmt.Sprintf("%s.plan.tfout", terraformDataDir)
}

func vars(vs map[string]string) []string {
	var args []string
	for k, v := range vs {
		args = append(args, "-var", fmt.Sprintf("%s=%s", k, v))
	}
	return args
}

func varFiles(vfs []string) []string {
	var args []string
	for _, v := range vfs {
		args = append(args, fmt.Sprintf("-var-file=%s", v))
	}
	return args
}

// helper function to write a netrc file.
// The following code comes from the official Git plugin for Drone:
// https://github.com/drone-plugins/drone-git/blob/8386effd2fe8c8695cf979427f8e1762bd805192/utils.go#L43-L68
func writeNetrc(machine, login, password string) error {
	if machine == "" {
		return nil
	}
	out := fmt.Sprintf(
		netrcFile,
		machine,
		login,
		password,
	)

	home := "/root"
	u, err := user.Current()
	if err == nil {
		home = u.HomeDir
	}
	path := filepath.Join(home, ".netrc")
	return ioutil.WriteFile(path, []byte(out), 0600)
}

const netrcFile = `
machine %s
login %s
password %s
`
<commit_msg>fix: -force is deprecated, use -auto-approve to delete in tf15<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sts"
)

type (
	// Config holds input parameters for the plugin
	Config struct {
		Actions          []string
		Vars             map[string]string
		Secrets          map[string]string
		InitOptions      InitOptions
		FmtOptions       FmtOptions
		Cacert           string
		Sensitive        bool
		RoleARN          string
		RootDir          string
		Parallelism      int
		Targets          []string
		VarFiles         []string
		TerraformDataDir string
		DisableRefresh   bool
	}

	// Netrc is credentials for cloning
	Netrc struct {
		Machine  string
		Login    string
		Password string
	}

	// InitOptions include options for the Terraform's init command
	InitOptions struct {
		BackendConfig []string `json:"backend-config"`
		Lock          *bool    `json:"lock"`
		LockTimeout   string   `json:"lock-timeout"`
	}

	// FmtOptions fmt options for the Terraform's fmt command
	FmtOptions struct {
		List  *bool `json:"list"`
		Write *bool `json:"write"`
		Diff  *bool `json:"diff"`
		Check *bool `json:"check"`
	}

	// Plugin represents the plugin instance to be executed
	Plugin struct {
		Config    Config
		Netrc     Netrc
		Terraform Terraform
	}
)

// Exec executes the plugin
func (p Plugin) Exec() error {
	// Install specified version of terraform
	if p.Terraform.Version != "" {
		err := installTerraform(p.Terraform.Version)

		if err != nil {
			return err
		}
	}

	if p.Config.RoleARN != "" && !credsSet() {
		assumeRole(p.Config.RoleARN)
	}

	// writing the .netrc file with Github credentials in it.
	err := writeNetrc(p.Netrc.Machine, p.Netrc.Login, p.Netrc.Password)
	if err != nil {
		return err
	}

	terraformDataDir := ".terraform"
	if p.Config.TerraformDataDir != "" {
		terraformDataDir = p.Config.TerraformDataDir
		os.Setenv("TF_DATA_DIR", p.Config.TerraformDataDir)
	}

	var commands []*exec.Cmd

	commands = append(commands, exec.Command("terraform", "version"))

	CopyTfEnv()

	if p.Config.Cacert != "" {
		commands = append(commands, installCaCert(p.Config.Cacert))
	}

	commands = append(commands, deleteCache(terraformDataDir))
	commands = append(commands, initCommand(p.Config.InitOptions))
	commands = append(commands, getModules())

	// Add commands listed from Actions
	for _, action := range p.Config.Actions {
		switch action {
		case "fmt":
			commands = append(commands, tfFmt(p.Config))
		case "validate":
			commands = append(commands, tfValidate())
		case "plan":
			commands = append(commands, tfPlan(p.Config, false))
		case "plan-destroy":
			commands = append(commands, tfPlan(p.Config, true))
		case "apply":
			commands = append(commands, tfApply(p.Config))
		case "destroy":
			commands = append(commands, tfDestroy(p.Config))
		default:
			return fmt.Errorf("valid actions are: fmt, validate, plan, apply, plan-destroy, destroy.  You provided %s", action)
		}
	}

	commands = append(commands, deleteCache(terraformDataDir))

	for _, c := range commands {
		if c.Dir == "" {
			wd, err := os.Getwd()
			if err == nil {
				c.Dir = wd
			}
		}
		if p.Config.RootDir != "" {
			c.Dir = c.Dir + "/" + p.Config.RootDir
		}
		c.Stdout = os.Stdout
		c.Stderr = os.Stderr
		if !p.Config.Sensitive {
			trace(c)
		}

		err := c.Run()
		if err != nil {
			logrus.WithFields(logrus.Fields{
				"error": err,
			}).Fatal("Failed to execute a command")
		}
		logrus.Debug("Command completed successfully")
	}

	return nil
}

// CopyTfEnv creates copies of TF_VAR_ to lowercase
func CopyTfEnv() {
	tfVar := regexp.MustCompile(`^TF_VAR_.*$`)
	for _, e := range os.Environ() {
		pair := strings.SplitN(e, "=", 2)
		if tfVar.MatchString(pair[0]) {
			name := strings.Split(pair[0], "TF_VAR_")
			os.Setenv(fmt.Sprintf("TF_VAR_%s", strings.ToLower(name[1])), pair[1])
		}
	}
}

func credsSet() bool {
	awsTokens := []string{"AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN"}
	for _, token := range awsTokens {
		if os.Getenv(token) == "" {
			return false
		}
	}
	return true
}

func assumeRole(roleArn string) {
	client := sts.New(session.New())
	duration := time.Hour * 1
	stsProvider := &stscreds.AssumeRoleProvider{
		Client:          client,
		Duration:        duration,
		RoleARN:         roleArn,
		RoleSessionName: "drone",
	}

	value, err := credentials.NewCredentials(stsProvider).Get()
	if err != nil {
		logrus.WithFields(logrus.Fields{
			"error": err,
		}).Fatal("Error assuming role!")
	}
	os.Setenv("AWS_ACCESS_KEY_ID", value.AccessKeyID)
	os.Setenv("AWS_SECRET_ACCESS_KEY", value.SecretAccessKey)
	os.Setenv("AWS_SESSION_TOKEN", value.SessionToken)
}

func deleteCache(terraformDataDir string) *exec.Cmd {
	return exec.Command(
		"rm",
		"-rf",
		terraformDataDir,
	)
}

func getModules() *exec.Cmd {
	return exec.Command(
		"terraform",
		"get",
	)
}

func initCommand(config InitOptions) *exec.Cmd {
	args := []string{
		"init",
	}

	for _, v := range config.BackendConfig {
		args = append(args, fmt.Sprintf("-backend-config=%s", v))
	}

	// True is default in TF
	if config.Lock != nil {
		args = append(args, fmt.Sprintf("-lock=%t", *config.Lock))
	}

	// "0s" is default in TF
	if config.LockTimeout != "" {
		args = append(args, fmt.Sprintf("-lock-timeout=%s", config.LockTimeout))
	}

	// Fail Terraform execution on prompt
	args = append(args, "-input=false")

	return exec.Command(
		"terraform",
		args...,
	)
}

func installCaCert(cacert string) *exec.Cmd {
	ioutil.WriteFile("/usr/local/share/ca-certificates/ca_cert.crt", []byte(cacert), 0644)
	return exec.Command(
		"update-ca-certificates",
	)
}

func trace(cmd *exec.Cmd) {
	fmt.Println("$", strings.Join(cmd.Args, " "))
}

func tfApply(config Config) *exec.Cmd {
	args := []string{
		"apply",
	}
	for _, v := range config.Targets {
		args = append(args, "--target", fmt.Sprintf("%s", v))
	}
	if config.Parallelism > 0 {
		args = append(args, fmt.Sprintf("-parallelism=%d", config.Parallelism))
	}
	if config.InitOptions.Lock != nil {
		args = append(args, fmt.Sprintf("-lock=%t", *config.InitOptions.Lock))
	}
	if config.InitOptions.LockTimeout != "" {
		args = append(args, fmt.Sprintf("-lock-timeout=%s", config.InitOptions.LockTimeout))
	}
	if config.DisableRefresh {
		args = append(args, "-refresh=false")
	}
	args = append(args, getTfoutPath())

	return exec.Command(
		"terraform",
		args...,
	)
}

func tfDestroy(config Config) *exec.Cmd {
	args := []string{
		"destroy",
	}
	for _, v := range config.Targets {
		args = append(args, fmt.Sprintf("-target=%s", v))
	}
	args = append(args, varFiles(config.VarFiles)...)
	args = append(args, vars(config.Vars)...)
	if config.Parallelism > 0 {
		args = append(args, fmt.Sprintf("-parallelism=%d", config.Parallelism))
	}
	if config.InitOptions.Lock != nil {
		args = append(args, fmt.Sprintf("-lock=%t", *config.InitOptions.Lock))
	}
	if config.InitOptions.LockTimeout != "" {
		args = append(args, fmt.Sprintf("-lock-timeout=%s", config.InitOptions.LockTimeout))
	}
	args = append(args, "-auto-approve")
	return exec.Command(
		"terraform",
		args...,
	)
}

func tfPlan(config Config, destroy bool) *exec.Cmd {
	args := []string{
		"plan",
	}

	if destroy {
		args = append(args, "-destroy")
	} else {
		args = append(args, fmt.Sprintf("-out=%s", getTfoutPath()))
	}

	for _, v := range config.Targets {
		args = append(args, "--target", fmt.Sprintf("%s", v))
	}
	args = append(args, varFiles(config.VarFiles)...)
	args = append(args, vars(config.Vars)...)
	if config.Parallelism > 0 {
		args = append(args, fmt.Sprintf("-parallelism=%d", config.Parallelism))
	}
	if config.InitOptions.Lock != nil {
		args = append(args, fmt.Sprintf("-lock=%t", *config.InitOptions.Lock))
	}
	if config.InitOptions.LockTimeout != "" {
		args = append(args, fmt.Sprintf("-lock-timeout=%s", config.InitOptions.LockTimeout))
	}
	if config.DisableRefresh {
		args = append(args, "-refresh=false")
	}
	return exec.Command(
		"terraform",
		args...,
	)
}

func tfValidate() *exec.Cmd {
	return exec.Command(
		"terraform",
		"validate",
	)
}

func tfFmt(config Config) *exec.Cmd {
	args := []string{
		"fmt",
	}
	if config.FmtOptions.List != nil {
		args = append(args, fmt.Sprintf("-list=%t", *config.FmtOptions.List))
	}
	if config.FmtOptions.Write != nil {
		args = append(args, fmt.Sprintf("-write=%t", *config.FmtOptions.Write))
	}
	if config.FmtOptions.Diff != nil {
		args = append(args, fmt.Sprintf("-diff=%t", *config.FmtOptions.Diff))
	}
	if config.FmtOptions.Check != nil {
		args = append(args, fmt.Sprintf("-check=%t", *config.FmtOptions.Check))
	}
	return exec.Command(
		"terraform",
		args...,
	)
}

func getTfoutPath() string {
	terraformDataDir := os.Getenv("TF_DATA_DIR")
	if terraformDataDir == ".terraform" || terraformDataDir == "" {
		return "plan.tfout"
	}

	return fmt.Sprintf("%s.plan.tfout", terraformDataDir)
}

func vars(vs map[string]string) []string {
	var args []string
	for k, v := range vs {
		args = append(args, "-var", fmt.Sprintf("%s=%s", k, v))
	}
	return args
}

func varFiles(vfs []string) []string {
	var args []string
	for _, v := range vfs {
		args = append(args, fmt.Sprintf("-var-file=%s", v))
	}
	return args
}

// helper function to write a netrc file.
// The following code comes from the official Git plugin for Drone:
// https://github.com/drone-plugins/drone-git/blob/8386effd2fe8c8695cf979427f8e1762bd805192/utils.go#L43-L68
func writeNetrc(machine, login, password string) error {
	if machine == "" {
		return nil
	}
	out := fmt.Sprintf(
		netrcFile,
		machine,
		login,
		password,
	)

	home := "/root"
	u, err := user.Current()
	if err == nil {
		home = u.HomeDir
	}
	path := filepath.Join(home, ".netrc")
	return ioutil.WriteFile(path, []byte(out), 0600)
}

const netrcFile = `
machine %s
login %s
password %s
`
<|endoftext|>


<commit_before>package main

import (
	"crypto/md5"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"strconv"

	_ "github.com/mattn/go-sqlite3"
)

const listenAddr = "127.0.0.1:9292"

var (
	db         *sql.DB
	getSetStmt *sql.Stmt
)

type Set struct {
	Id             int
	ThumbPhotoId   int
	Name           string
	PhotosCount    int
	TakenAt        sql.NullString
	ThumbPhotoPath string
}

func (s *Set) ThumbURL() string {
	identifier := fmt.Sprintf("%x", md5.Sum([]byte(s.ThumbPhotoPath)))
	return fmt.Sprintf("/thumbs/%s_small.jpg", identifier)
}

func (s *Set) MarshalJSON() ([]byte, error) { // implements Marshaler
	setMap := map[string]interface{}{
		"id":             s.Id,
		"thumb_photo_id": s.ThumbPhotoId,
		"name":           s.Name,
		"photos_count":   s.PhotosCount,
		"thumb_url":      s.ThumbURL(),
	}
	setMap["taken_at"], _ = s.TakenAt.Value()
	return json.Marshal(setMap)
}

func badRequest(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusBadRequest)
	fmt.Fprintln(w, http.StatusBadRequest, "Bad Request")
}

func requireParam(param string, w http.ResponseWriter, r *http.Request) error {
	if len(r.URL.Query()[param]) == 0 {
		badRequest(w, r)
		return errors.New(fmt.Sprintf("missing %s parameter", param))
	}
	return nil
}

func getSetById(setId int) (*Set, error) {
	set := Set{}
	row := getSetStmt.QueryRow(setId)
	err := row.Scan(
		&set.Id,
		&set.ThumbPhotoId,
		&set.Name,
		&set.PhotosCount,
		&set.TakenAt,
		&set.ThumbPhotoPath,
	)
	return &set, err
}

func getSetHandler(w http.ResponseWriter, r *http.Request) {
	if requireParam("id", w, r) != nil {
		return
	}

	setId, err := strconv.Atoi(r.URL.Query()["id"][0])
	if err != nil {
		log.Fatalln("Failed to convert id to integer:", err)
	}

	set, err := getSetById(setId)
	if err == sql.ErrNoRows { // set does not exist
		http.NotFound(w, r)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(set)
}

func main() {
	db, err := sql.Open("sqlite3", "thyme.db")
	if err != nil {
		log.Fatalln("Failed to open database:", err)
	}
	defer db.Close()

	getSetStmt, err = db.Prepare(`
	SELECT sets.id, thumb_photo_id, name, photos_count, sets.taken_at, photos.path
	FROM sets JOIN photos ON sets.thumb_photo_id = photos.id
	WHERE sets.id = ?
	`)
	if err != nil {
		log.Fatalln("Failed to access table:", err)
	}
	defer getSetStmt.Close()

	http.Handle("/", http.FileServer(http.Dir("public"))) // static
	http.HandleFunc("/set", getSetHandler)

	fmt.Printf("Listening on http://%s\n", listenAddr)
	fmt.Println("Press Ctrl-C to exit")

	log.Fatal(http.ListenAndServe(listenAddr, nil))
}
<commit_msg>Sort set columns alphabetically<commit_after>package main

import (
	"crypto/md5"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"strconv"

	_ "github.com/mattn/go-sqlite3"
)

const listenAddr = "127.0.0.1:9292"

var (
	db         *sql.DB
	getSetStmt *sql.Stmt
)

type Set struct {
	Id             int
	Name           string
	PhotosCount    int
	TakenAt        sql.NullString
	ThumbPhotoId   int
	ThumbPhotoPath string
}

func (s *Set) ThumbURL() string {
	identifier := fmt.Sprintf("%x", md5.Sum([]byte(s.ThumbPhotoPath)))
	return fmt.Sprintf("/thumbs/%s_small.jpg", identifier)
}

func (s *Set) MarshalJSON() ([]byte, error) { // implements Marshaler
	setMap := map[string]interface{}{
		"id":             s.Id,
		"name":           s.Name,
		"photos_count":   s.PhotosCount,
		"thumb_photo_id": s.ThumbPhotoId,
		"thumb_url":      s.ThumbURL(),
	}
	setMap["taken_at"], _ = s.TakenAt.Value()
	return json.Marshal(setMap)
}

func badRequest(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusBadRequest)
	fmt.Fprintln(w, http.StatusBadRequest, "Bad Request")
}

func requireParam(param string, w http.ResponseWriter, r *http.Request) error {
	if len(r.URL.Query()[param]) == 0 {
		badRequest(w, r)
		return errors.New(fmt.Sprintf("missing %s parameter", param))
	}
	return nil
}

func getSetById(setId int) (*Set, error) {
	set := Set{}
	row := getSetStmt.QueryRow(setId)
	err := row.Scan(
		&set.Id,
		&set.Name,
		&set.PhotosCount,
		&set.TakenAt,
		&set.ThumbPhotoId,
		&set.ThumbPhotoPath,
	)
	return &set, err
}

func getSetHandler(w http.ResponseWriter, r *http.Request) {
	if requireParam("id", w, r) != nil {
		return
	}

	setId, err := strconv.Atoi(r.URL.Query()["id"][0])
	if err != nil {
		log.Fatalln("Failed to convert id to integer:", err)
	}

	set, err := getSetById(setId)
	if err == sql.ErrNoRows { // set does not exist
		http.NotFound(w, r)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(set)
}

func main() {
	db, err := sql.Open("sqlite3", "thyme.db")
	if err != nil {
		log.Fatalln("Failed to open database:", err)
	}
	defer db.Close()

	getSetStmt, err = db.Prepare(`
	SELECT
	sets.id, name, photos_count, sets.taken_at, thumb_photo_id, photos.path
	FROM sets
	JOIN photos ON sets.thumb_photo_id = photos.id
	WHERE sets.id = ?
	`)
	if err != nil {
		log.Fatalln("Failed to access table:", err)
	}
	defer getSetStmt.Close()

	http.Handle("/", http.FileServer(http.Dir("public"))) // static
	http.HandleFunc("/set", getSetHandler)

	fmt.Printf("Listening on http://%s\n", listenAddr)
	fmt.Println("Press Ctrl-C to exit")

	log.Fatal(http.ListenAndServe(listenAddr, nil))
}
<|endoftext|>


<commit_before>// Copyright 2015 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"flag"
	"fmt"
	"net"
	"os"
	"os/signal"
	"runtime"
	"syscall"
	"time"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/ngaut/systimemon"
	"github.com/pingcap/tidb"
	"github.com/pingcap/tidb/ddl"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/perfschema"
	"github.com/pingcap/tidb/plan"
	"github.com/pingcap/tidb/privilege/privileges"
	"github.com/pingcap/tidb/server"
	"github.com/pingcap/tidb/sessionctx/binloginfo"
	"github.com/pingcap/tidb/store/localstore/boltdb"
	"github.com/pingcap/tidb/store/tikv"
	"github.com/pingcap/tidb/util/printer"
	"github.com/pingcap/tipb/go-binlog"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/push"
	"google.golang.org/grpc"
)

var (
	version         = flag.Bool("V", false, "print version information and exit")
	store           = flag.String("store", "goleveldb", "registered store name, [memory, goleveldb, boltdb, tikv, mocktikv]")
	storePath       = flag.String("path", "/tmp/tidb", "tidb storage path")
	logLevel        = flag.String("L", "info", "log level: info, debug, warn, error, fatal")
	host            = flag.String("host", "0.0.0.0", "tidb server host")
	port            = flag.String("P", "4000", "tidb server port")
	statusPort      = flag.String("status", "10080", "tidb server status port")
	lease           = flag.String("lease", "1s", "schema lease duration, very dangerous to change only if you know what you do")
	socket          = flag.String("socket", "", "The socket file to use for connection.")
	enablePS        = flag.Bool("perfschema", false, "If enable performance schema.")
	enablePrivilege = flag.Bool("privilege", true, "If enable privilege check feature. This flag will be removed in the future.")
	reportStatus    = flag.Bool("report-status", true, "If enable status report HTTP service.")
	logFile         = flag.String("log-file", "", "log file path")
	joinCon         = flag.Int("join-concurrency", 5, "the number of goroutines that participate joining.")
	crossJoin       = flag.Bool("cross-join", true, "whether support cartesian product or not.")
	metricsAddr     = flag.String("metrics-addr", "", "prometheus pushgateway address, leaves it empty will disable prometheus push.")
	metricsInterval = flag.Int("metrics-interval", 15, "prometheus client push interval in second, set \"0\" to disable prometheus push.")
	binlogSocket    = flag.String("binlog-socket", "", "socket file to write binlog")
	runDDL          = flag.Bool("run-ddl", true, "run ddl worker on this tidb-server")
	retryLimit      = flag.Int("retry-limit", 10, "the maximum number of retries when commit a transaction")
	skipGrantTable  = flag.Bool("skip-grant-table", false, "This option causes the server to start without using the privilege system at all.")

	timeJumpBackCounter = prometheus.NewCounter(
		prometheus.CounterOpts{
			Namespace: "tidb",
			Subsystem: "monitor",
			Name:      "time_jump_back_total",
			Help:      "Counter of system time jumps backward.",
		})
)

func main() {
	tidb.RegisterLocalStore("boltdb", boltdb.Driver{})
	tidb.RegisterStore("tikv", tikv.Driver{})
	tidb.RegisterStore("mocktikv", tikv.MockDriver{})

	runtime.GOMAXPROCS(runtime.NumCPU())

	flag.Parse()
	if *version {
		printer.PrintRawTiDBInfo()
		os.Exit(0)
	}
	if *skipGrantTable && !hasRootPrivilege() {
		log.Error("TiDB run with skip-grant-table need root privilege.")
		os.Exit(-1)
	}

	leaseDuration := parseLease()
	tidb.SetSchemaLease(leaseDuration)
	ddl.RunWorker = *runDDL
	tidb.SetCommitRetryLimit(*retryLimit)

	cfg := &server.Config{
		Addr:         fmt.Sprintf("%s:%s", *host, *port),
		LogLevel:     *logLevel,
		StatusAddr:   fmt.Sprintf(":%s", *statusPort),
		Socket:       *socket,
		ReportStatus: *reportStatus,
		Store:        *store,
		StorePath:    *storePath,
	}

	// set log options
	if len(*logFile) > 0 {
		err := log.SetOutputByName(*logFile)
		if err != nil {
			log.Fatal(errors.ErrorStack(err))
		}
		log.SetRotateByDay()
		log.SetHighlighting(false)
	}

	if joinCon != nil && *joinCon > 0 {
		plan.JoinConcurrency = *joinCon
	}
	plan.AllowCartesianProduct = *crossJoin
	// Call this before setting log level to make sure that TiDB info could be printed.
	printer.PrintTiDBInfo()
	log.SetLevelByString(cfg.LogLevel)

	store := createStore()

	if *enablePS {
		perfschema.EnablePerfSchema()
	}
	privileges.Enable = *enablePrivilege
	privileges.SkipWithGrant = *skipGrantTable
	if *binlogSocket != "" {
		createBinlogClient()
	}

	// Bootstrap a session to load information schema.
	_, err := tidb.BootstrapSession(store)
	if err != nil {
		log.Fatal(errors.ErrorStack(err))
	}

	var driver server.IDriver
	driver = server.NewTiDBDriver(store)
	var svr *server.Server
	svr, err = server.NewServer(cfg, driver)
	if err != nil {
		log.Fatal(errors.ErrorStack(err))
	}

	sc := make(chan os.Signal, 1)
	signal.Notify(sc,
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT)

	go func() {
		sig := <-sc
		log.Infof("Got signal [%d] to exit.", sig)
		svr.Close()
		os.Exit(0)
	}()

	prometheus.MustRegister(timeJumpBackCounter)
	go systimemon.StartMonitor(time.Now, func() {
		timeJumpBackCounter.Inc()
	})

	pushMetric(*metricsAddr, time.Duration(*metricsInterval)*time.Second)

	log.Error(svr.Run())
}

func createStore() kv.Storage {
	fullPath := fmt.Sprintf("%s://%s", *store, *storePath)
	store, err := tidb.NewStore(fullPath)
	if err != nil {
		log.Fatal(errors.ErrorStack(err))
	}
	return store
}

func createBinlogClient() {
	dialerOpt := grpc.WithDialer(func(addr string, timeout time.Duration) (net.Conn, error) {
		return net.DialTimeout("unix", addr, timeout)
	})
	clientCon, err := grpc.Dial(*binlogSocket, dialerOpt, grpc.WithInsecure())
	if err != nil {
		log.Fatal(errors.ErrorStack(err))
	}
	binloginfo.PumpClient = binlog.NewPumpClient(clientCon)
	log.Infof("created binlog client at %s", *binlogSocket)
}

// Prometheus push.
const zeroDuration = time.Duration(0)

// PushMetric pushs metircs in background.
func pushMetric(addr string, interval time.Duration) {
	if interval == zeroDuration || len(addr) == 0 {
		log.Info("disable Prometheus push client")
		return
	}
	log.Infof("start Prometheus push client with server addr %s and interval %s", addr, interval)
	go prometheusPushClient(addr, interval)
}

// PrometheusPushClient pushs metrics to Prometheus Pushgateway.
func prometheusPushClient(addr string, interval time.Duration) {
	// TODO: TiDB do not have uniq name, so we use host+port to compose a name.
	job := "tidb"
	for {
		err := push.AddFromGatherer(
			job, push.HostnameGroupingKey(),
			addr,
			prometheus.DefaultGatherer,
		)
		if err != nil {
			log.Errorf("could not push metrics to Prometheus Pushgateway: %v", err)
		}
		time.Sleep(interval)
	}
}

// parseLease parses lease argument string.
func parseLease() time.Duration {
	dur, err := time.ParseDuration(*lease)
	if err != nil {
		dur, err = time.ParseDuration(*lease + "s")
	}
	if err != nil || dur < 0 {
		log.Fatalf("invalid lease duration %s", *lease)
	}
	return dur
}

func hasRootPrivilege() bool {
	return os.Geteuid() == 0
}
<commit_msg>tidb-server: add domain close,adjust exit position (#3349)<commit_after>// Copyright 2015 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"flag"
	"fmt"
	"net"
	"os"
	"os/signal"
	"runtime"
	"syscall"
	"time"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/ngaut/systimemon"
	"github.com/pingcap/tidb"
	"github.com/pingcap/tidb/ddl"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/perfschema"
	"github.com/pingcap/tidb/plan"
	"github.com/pingcap/tidb/privilege/privileges"
	"github.com/pingcap/tidb/server"
	"github.com/pingcap/tidb/sessionctx/binloginfo"
	"github.com/pingcap/tidb/store/localstore/boltdb"
	"github.com/pingcap/tidb/store/tikv"
	"github.com/pingcap/tidb/util/printer"
	"github.com/pingcap/tipb/go-binlog"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/push"
	"google.golang.org/grpc"
)

var (
	version         = flag.Bool("V", false, "print version information and exit")
	store           = flag.String("store", "goleveldb", "registered store name, [memory, goleveldb, boltdb, tikv, mocktikv]")
	storePath       = flag.String("path", "/tmp/tidb", "tidb storage path")
	logLevel        = flag.String("L", "info", "log level: info, debug, warn, error, fatal")
	host            = flag.String("host", "0.0.0.0", "tidb server host")
	port            = flag.String("P", "4000", "tidb server port")
	statusPort      = flag.String("status", "10080", "tidb server status port")
	lease           = flag.String("lease", "1s", "schema lease duration, very dangerous to change only if you know what you do")
	socket          = flag.String("socket", "", "The socket file to use for connection.")
	enablePS        = flag.Bool("perfschema", false, "If enable performance schema.")
	enablePrivilege = flag.Bool("privilege", true, "If enable privilege check feature. This flag will be removed in the future.")
	reportStatus    = flag.Bool("report-status", true, "If enable status report HTTP service.")
	logFile         = flag.String("log-file", "", "log file path")
	joinCon         = flag.Int("join-concurrency", 5, "the number of goroutines that participate joining.")
	crossJoin       = flag.Bool("cross-join", true, "whether support cartesian product or not.")
	metricsAddr     = flag.String("metrics-addr", "", "prometheus pushgateway address, leaves it empty will disable prometheus push.")
	metricsInterval = flag.Int("metrics-interval", 15, "prometheus client push interval in second, set \"0\" to disable prometheus push.")
	binlogSocket    = flag.String("binlog-socket", "", "socket file to write binlog")
	runDDL          = flag.Bool("run-ddl", true, "run ddl worker on this tidb-server")
	retryLimit      = flag.Int("retry-limit", 10, "the maximum number of retries when commit a transaction")
	skipGrantTable  = flag.Bool("skip-grant-table", false, "This option causes the server to start without using the privilege system at all.")

	timeJumpBackCounter = prometheus.NewCounter(
		prometheus.CounterOpts{
			Namespace: "tidb",
			Subsystem: "monitor",
			Name:      "time_jump_back_total",
			Help:      "Counter of system time jumps backward.",
		})
)

func main() {
	tidb.RegisterLocalStore("boltdb", boltdb.Driver{})
	tidb.RegisterStore("tikv", tikv.Driver{})
	tidb.RegisterStore("mocktikv", tikv.MockDriver{})

	runtime.GOMAXPROCS(runtime.NumCPU())

	flag.Parse()
	if *version {
		printer.PrintRawTiDBInfo()
		os.Exit(0)
	}
	if *skipGrantTable && !hasRootPrivilege() {
		log.Error("TiDB run with skip-grant-table need root privilege.")
		os.Exit(-1)
	}

	leaseDuration := parseLease()
	tidb.SetSchemaLease(leaseDuration)
	ddl.RunWorker = *runDDL
	tidb.SetCommitRetryLimit(*retryLimit)

	cfg := &server.Config{
		Addr:         fmt.Sprintf("%s:%s", *host, *port),
		LogLevel:     *logLevel,
		StatusAddr:   fmt.Sprintf(":%s", *statusPort),
		Socket:       *socket,
		ReportStatus: *reportStatus,
		Store:        *store,
		StorePath:    *storePath,
	}

	// set log options
	if len(*logFile) > 0 {
		err := log.SetOutputByName(*logFile)
		if err != nil {
			log.Fatal(errors.ErrorStack(err))
		}
		log.SetRotateByDay()
		log.SetHighlighting(false)
	}

	if joinCon != nil && *joinCon > 0 {
		plan.JoinConcurrency = *joinCon
	}
	plan.AllowCartesianProduct = *crossJoin
	// Call this before setting log level to make sure that TiDB info could be printed.
	printer.PrintTiDBInfo()
	log.SetLevelByString(cfg.LogLevel)

	store := createStore()

	if *enablePS {
		perfschema.EnablePerfSchema()
	}
	privileges.Enable = *enablePrivilege
	privileges.SkipWithGrant = *skipGrantTable
	if *binlogSocket != "" {
		createBinlogClient()
	}

	// Bootstrap a session to load information schema.
	domain, err := tidb.BootstrapSession(store)
	if err != nil {
		log.Fatal(errors.ErrorStack(err))
	}

	var driver server.IDriver
	driver = server.NewTiDBDriver(store)
	var svr *server.Server
	svr, err = server.NewServer(cfg, driver)
	if err != nil {
		log.Fatal(errors.ErrorStack(err))
	}

	sc := make(chan os.Signal, 1)
	signal.Notify(sc,
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT)

	go func() {
		sig := <-sc
		log.Infof("Got signal [%d] to exit.", sig)
		svr.Close()
	}()

	prometheus.MustRegister(timeJumpBackCounter)
	go systimemon.StartMonitor(time.Now, func() {
		timeJumpBackCounter.Inc()
	})

	pushMetric(*metricsAddr, time.Duration(*metricsInterval)*time.Second)

	log.Error(svr.Run())
	domain.Close()
	os.Exit(0)
}

func createStore() kv.Storage {
	fullPath := fmt.Sprintf("%s://%s", *store, *storePath)
	store, err := tidb.NewStore(fullPath)
	if err != nil {
		log.Fatal(errors.ErrorStack(err))
	}
	return store
}

func createBinlogClient() {
	dialerOpt := grpc.WithDialer(func(addr string, timeout time.Duration) (net.Conn, error) {
		return net.DialTimeout("unix", addr, timeout)
	})
	clientCon, err := grpc.Dial(*binlogSocket, dialerOpt, grpc.WithInsecure())
	if err != nil {
		log.Fatal(errors.ErrorStack(err))
	}
	binloginfo.PumpClient = binlog.NewPumpClient(clientCon)
	log.Infof("created binlog client at %s", *binlogSocket)
}

// Prometheus push.
const zeroDuration = time.Duration(0)

// PushMetric pushs metircs in background.
func pushMetric(addr string, interval time.Duration) {
	if interval == zeroDuration || len(addr) == 0 {
		log.Info("disable Prometheus push client")
		return
	}
	log.Infof("start Prometheus push client with server addr %s and interval %s", addr, interval)
	go prometheusPushClient(addr, interval)
}

// PrometheusPushClient pushs metrics to Prometheus Pushgateway.
func prometheusPushClient(addr string, interval time.Duration) {
	// TODO: TiDB do not have uniq name, so we use host+port to compose a name.
	job := "tidb"
	for {
		err := push.AddFromGatherer(
			job, push.HostnameGroupingKey(),
			addr,
			prometheus.DefaultGatherer,
		)
		if err != nil {
			log.Errorf("could not push metrics to Prometheus Pushgateway: %v", err)
		}
		time.Sleep(interval)
	}
}

// parseLease parses lease argument string.
func parseLease() time.Duration {
	dur, err := time.ParseDuration(*lease)
	if err != nil {
		dur, err = time.ParseDuration(*lease + "s")
	}
	if err != nil || dur < 0 {
		log.Fatalf("invalid lease duration %s", *lease)
	}
	return dur
}

func hasRootPrivilege() bool {
	return os.Geteuid() == 0
}
<|endoftext|>


<commit_before>// Command aws-gen-gocli parses a JSON description of an AWS API and generates a
// Go file containing a client for the API.
//
//     aws-gen-gocli apis/s3/2006-03-03.normal.json
package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime/debug"
	"sort"
	"strings"
	"sync"

	"github.com/aws/aws-sdk-go/internal/model/api"
	"github.com/aws/aws-sdk-go/internal/util"
)

type generateInfo struct {
	*api.API
	PackageDir string
}

// newGenerateInfo initializes the service API's folder structure for a specific service.
// If the SERVICES environment variable is set, and this service is not apart of the list
// this service will be skipped.
func newGenerateInfo(modelFile, svcPath string) *generateInfo {
	g := &generateInfo{API: &api.API{}}
	g.API.Attach(modelFile)

	paginatorsFile := strings.Replace(modelFile, ".normal.json", ".paginators.json", -1)
	if _, err := os.Stat(paginatorsFile); err == nil {
		g.API.AttachPaginators(paginatorsFile)
	}

	if svc := os.Getenv("SERVICES"); svc != "" {
		svcs := strings.Split(svc, ",")
		for _, s := range svcs {
			if s != g.API.PackageName() { // skip this non-included service
				return nil
			}
		}
	}

	// ensure the directory exists
	pkgDir := filepath.Join(svcPath, g.API.PackageName())
	os.MkdirAll(pkgDir, 0775)
	os.MkdirAll(filepath.Join(pkgDir, g.API.InterfacePackageName()), 0775)

	g.PackageDir = pkgDir

	return g
}

// Generates service api, examples, and interface from api json definition files.
//
// Flags:
// -path alternative service path to write generated files to for each service.
//
// Env:
//  SERVICES comma separated list of services to generate.
func main() {
	var svcPath string
	flag.StringVar(&svcPath, "path", "service", "generate in a specific directory (default: 'service')")
	flag.Parse()

	files := []string{}
	for i := 0; i < flag.NArg(); i++ {
		file := flag.Arg(i)
		if strings.Contains(file, "*") {
			paths, _ := filepath.Glob(file)
			files = append(files, paths...)
		} else {
			files = append(files, file)
		}
	}

	sort.Strings(files)

	// Remove old API versions from list
	m := map[string]bool{}
	for i := range files {
		idx := len(files) - 1 - i
		parts := strings.Split(files[idx], string(filepath.Separator))
		svc := parts[len(parts)-2] // service name is 2nd-to-last component

		if m[svc] {
			files[idx] = "" // wipe this one out if we already saw the service
		}
		m[svc] = true
	}

	w := sync.WaitGroup{}
	for i := range files {
		file := files[i]
		if file == "" { // empty file
			continue
		}

		w.Add(1)
		go func() {
			defer func() {
				w.Done()
				if r := recover(); r != nil {
					fmtStr := "Error generating %s\n%s\n%s\n"
					fmt.Fprintf(os.Stderr, fmtStr, file, r, debug.Stack())
				}
			}()

			if g := newGenerateInfo(file, svcPath); g != nil {
				switch g.API.PackageName() {
				case "simpledb", "importexport":
					// These services are not yet supported, do nothing.
				default:
					fmt.Printf("Generating %s (%s)...\n",
						g.API.PackageName(), g.API.Metadata.APIVersion)

					// write api.go and service.go files
					g.writeAPIFile()
					g.writeExamplesFile()
					g.writeServiceFile()
					g.writeInterfaceFile()
				}
			}
		}()
	}
	w.Wait()
}

const codeLayout = `// THIS FILE IS AUTOMATICALLY GENERATED. DO NOT EDIT.
%s
package %s

%s
`

func writeGoFile(file string, layout string, args ...interface{}) error {
	return ioutil.WriteFile(file, []byte(util.GoFmt(fmt.Sprintf(layout, args...))), 0664)
}

// writeExamplesFile writes out the service example file.
func (g *generateInfo) writeExamplesFile() {
	writeGoFile(filepath.Join(g.PackageDir, "examples_test.go"),
		codeLayout,
		"",
		g.API.PackageName()+"_test",
		g.API.ExampleGoCode(),
	)
}

// writeServiceFile writes out the service initialization file.
func (g *generateInfo) writeServiceFile() {
	writeGoFile(filepath.Join(g.PackageDir, "service.go"),
		codeLayout,
		"",
		g.API.PackageName(),
		g.API.ServiceGoCode(),
	)
}

// writeInterfaceFile writes out the service interface file.
func (g *generateInfo) writeInterfaceFile() {
	writeGoFile(filepath.Join(g.PackageDir, g.API.InterfacePackageName(), "interface.go"),
		codeLayout,
		fmt.Sprintf("\n// Package %s provides an interface for the %s.",
			g.API.InterfacePackageName(), g.API.Metadata.ServiceFullName),
		g.API.InterfacePackageName(),
		g.API.InterfaceGoCode(),
	)

	writeGoFile(filepath.Join(g.PackageDir, g.API.InterfacePackageName(), "interface_test.go"),
		codeLayout,
		"",
		g.API.InterfacePackageName()+"_test",
		g.API.InterfaceTestGoCode(),
	)
}

// writeAPIFile writes out the service api file.
func (g *generateInfo) writeAPIFile() {
	writeGoFile(filepath.Join(g.PackageDir, "api.go"),
		codeLayout,
		fmt.Sprintf("\n// Package %s provides a client for %s.",
			g.API.PackageName(), g.API.Metadata.ServiceFullName),
		g.API.PackageName(),
		g.API.APIGoCode(),
	)
}
<commit_msg>Cleanedup API generation to not create unsupported service folders<commit_after>// Command aws-gen-gocli parses a JSON description of an AWS API and generates a
// Go file containing a client for the API.
//
//     aws-gen-gocli apis/s3/2006-03-03.normal.json
package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime/debug"
	"sort"
	"strings"
	"sync"

	"github.com/aws/aws-sdk-go/internal/model/api"
	"github.com/aws/aws-sdk-go/internal/util"
)

type generateInfo struct {
	*api.API
	PackageDir string
}

var excludeServices = map[string]struct{}{
	"simpledb":     struct{}{},
	"importexport": struct{}{},
}

// newGenerateInfo initializes the service API's folder structure for a specific service.
// If the SERVICES environment variable is set, and this service is not apart of the list
// this service will be skipped.
func newGenerateInfo(modelFile, svcPath string) *generateInfo {
	g := &generateInfo{API: &api.API{}}
	g.API.Attach(modelFile)

	if _, ok := excludeServices[g.API.PackageName()]; ok {
		return nil
	}

	paginatorsFile := strings.Replace(modelFile, ".normal.json", ".paginators.json", -1)
	if _, err := os.Stat(paginatorsFile); err == nil {
		g.API.AttachPaginators(paginatorsFile)
	}

	if svc := os.Getenv("SERVICES"); svc != "" {
		svcs := strings.Split(svc, ",")

		included := false
		for _, s := range svcs {
			if s == g.API.PackageName() {
				included = true
				break
			}
		}
		if !included {
			// skip this non-included service
			return nil
		}
	}

	// ensure the directory exists
	pkgDir := filepath.Join(svcPath, g.API.PackageName())
	os.MkdirAll(pkgDir, 0775)
	os.MkdirAll(filepath.Join(pkgDir, g.API.InterfacePackageName()), 0775)

	g.PackageDir = pkgDir

	return g
}

// Generates service api, examples, and interface from api json definition files.
//
// Flags:
// -path alternative service path to write generated files to for each service.
//
// Env:
//  SERVICES comma separated list of services to generate.
func main() {
	var svcPath string
	flag.StringVar(&svcPath, "path", "service", "generate in a specific directory (default: 'service')")
	flag.Parse()

	files := []string{}
	for i := 0; i < flag.NArg(); i++ {
		file := flag.Arg(i)
		if strings.Contains(file, "*") {
			paths, _ := filepath.Glob(file)
			files = append(files, paths...)
		} else {
			files = append(files, file)
		}
	}

	for svcName := range excludeServices {
		if strings.Contains(os.Getenv("SERVICES"), svcName) {
			fmt.Printf("Service %s is not supported\n", svcName)
			os.Exit(1)
		}
	}

	sort.Strings(files)

	// Remove old API versions from list
	m := map[string]bool{}
	for i := range files {
		idx := len(files) - 1 - i
		parts := strings.Split(files[idx], string(filepath.Separator))
		svc := parts[len(parts)-2] // service name is 2nd-to-last component

		if m[svc] {
			files[idx] = "" // wipe this one out if we already saw the service
		}
		m[svc] = true
	}

	w := sync.WaitGroup{}
	for i := range files {
		file := files[i]
		if file == "" { // empty file
			continue
		}

		w.Add(1)
		go func() {
			defer func() {
				w.Done()
				if r := recover(); r != nil {
					fmtStr := "Error generating %s\n%s\n%s\n"
					fmt.Fprintf(os.Stderr, fmtStr, file, r, debug.Stack())
				}
			}()

			if g := newGenerateInfo(file, svcPath); g != nil {
				if _, ok := excludeServices[g.API.PackageName()]; !ok {
					// Skip services not yet supported.

					fmt.Printf("Generating %s (%s)...\n",
						g.API.PackageName(), g.API.Metadata.APIVersion)

					// write api.go and service.go files
					g.writeAPIFile()
					g.writeExamplesFile()
					g.writeServiceFile()
					g.writeInterfaceFile()
				}
			}
		}()
	}
	w.Wait()
}

const codeLayout = `// THIS FILE IS AUTOMATICALLY GENERATED. DO NOT EDIT.
%s
package %s

%s
`

func writeGoFile(file string, layout string, args ...interface{}) error {
	return ioutil.WriteFile(file, []byte(util.GoFmt(fmt.Sprintf(layout, args...))), 0664)
}

// writeExamplesFile writes out the service example file.
func (g *generateInfo) writeExamplesFile() {
	writeGoFile(filepath.Join(g.PackageDir, "examples_test.go"),
		codeLayout,
		"",
		g.API.PackageName()+"_test",
		g.API.ExampleGoCode(),
	)
}

// writeServiceFile writes out the service initialization file.
func (g *generateInfo) writeServiceFile() {
	writeGoFile(filepath.Join(g.PackageDir, "service.go"),
		codeLayout,
		"",
		g.API.PackageName(),
		g.API.ServiceGoCode(),
	)
}

// writeInterfaceFile writes out the service interface file.
func (g *generateInfo) writeInterfaceFile() {
	writeGoFile(filepath.Join(g.PackageDir, g.API.InterfacePackageName(), "interface.go"),
		codeLayout,
		fmt.Sprintf("\n// Package %s provides an interface for the %s.",
			g.API.InterfacePackageName(), g.API.Metadata.ServiceFullName),
		g.API.InterfacePackageName(),
		g.API.InterfaceGoCode(),
	)

	writeGoFile(filepath.Join(g.PackageDir, g.API.InterfacePackageName(), "interface_test.go"),
		codeLayout,
		"",
		g.API.InterfacePackageName()+"_test",
		g.API.InterfaceTestGoCode(),
	)
}

// writeAPIFile writes out the service api file.
func (g *generateInfo) writeAPIFile() {
	writeGoFile(filepath.Join(g.PackageDir, "api.go"),
		codeLayout,
		fmt.Sprintf("\n// Package %s provides a client for %s.",
			g.API.PackageName(), g.API.Metadata.ServiceFullName),
		g.API.PackageName(),
		g.API.APIGoCode(),
	)
}
<|endoftext|>


<commit_before>// Copyright 2020 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package postgres

import (
	"context"
	"errors"
	"sort"
	"testing"

	"github.com/google/go-cmp/cmp"
	"github.com/google/go-cmp/cmp/cmpopts"
	"golang.org/x/pkgsite/internal"
	"golang.org/x/pkgsite/internal/licenses"
	"golang.org/x/pkgsite/internal/stdlib"
	"golang.org/x/pkgsite/internal/testing/sample"
)

func TestGetLicenses(t *testing.T) {
	testModule := sample.LegacyModule(sample.ModulePath, "v1.2.3", "A/B")
	stdlibModule := sample.LegacyModule(stdlib.ModulePath, "v1.13.0", "cmd/go")
	mit := &licenses.Metadata{Types: []string{"MIT"}, FilePath: "LICENSE"}
	bsd := &licenses.Metadata{Types: []string{"BSD-3-Clause"}, FilePath: "A/B/LICENSE"}

	mitLicense := &licenses.License{Metadata: mit}
	bsdLicense := &licenses.License{Metadata: bsd}
	testModule.Licenses = []*licenses.License{bsdLicense, mitLicense}
	sort.Slice(testModule.Units, func(i, j int) bool {
		return testModule.Units[i].Path < testModule.Units[j].Path
	})

	// github.com/valid/module_name
	testModule.Units[0].Licenses = []*licenses.Metadata{mit}
	// github.com/valid/module_name/A
	testModule.Units[1].Licenses = []*licenses.Metadata{mit}
	// github.com/valid/module_name/A/B
	testModule.Units[2].Licenses = []*licenses.Metadata{mit, bsd}

	defer ResetTestDB(testDB, t)
	ctx, cancel := context.WithTimeout(context.Background(), testTimeout*5)
	defer cancel()
	if err := testDB.InsertModule(ctx, testModule); err != nil {
		t.Fatal(err)
	}
	if err := testDB.InsertModule(ctx, stdlibModule); err != nil {
		t.Fatal(err)
	}
	for _, test := range []struct {
		err                                 error
		name, fullPath, modulePath, version string
		want                                []*licenses.License
	}{
		{
			name:       "module root",
			fullPath:   sample.ModulePath,
			modulePath: sample.ModulePath,
			version:    testModule.Version,
			want:       []*licenses.License{testModule.Licenses[1]},
		},
		{
			name:       "package without license",
			fullPath:   sample.ModulePath + "/A",
			modulePath: sample.ModulePath,
			version:    testModule.Version,
			want:       []*licenses.License{testModule.Licenses[1]},
		},
		{
			name:       "package with additional license",
			fullPath:   sample.ModulePath + "/A/B",
			modulePath: sample.ModulePath,
			version:    testModule.Version,
			want:       testModule.Licenses,
		},
		{
			name:       "stdlib directory",
			fullPath:   "cmd",
			modulePath: stdlib.ModulePath,
			version:    stdlibModule.Version,
			want:       stdlibModule.Licenses,
		},
		{
			name:       "stdlib package",
			fullPath:   "cmd/go",
			modulePath: stdlib.ModulePath,
			version:    stdlibModule.Version,
			want:       stdlibModule.Licenses,
		},
		{
			name:       "stdlib module",
			fullPath:   stdlib.ModulePath,
			modulePath: stdlib.ModulePath,
			version:    stdlibModule.Version,
			want:       stdlibModule.Licenses,
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			u, err := testDB.GetUnit(ctx, &internal.UnitMeta{Path: test.fullPath, ModulePath: test.modulePath, Version: test.version}, internal.WithLicenses)
			if !errors.Is(err, test.err) {
				t.Fatal(err)
			}
			got := u.LicenseContents
			sort.Slice(got, func(i, j int) bool {
				return got[i].FilePath < got[j].FilePath
			})
			sort.Slice(test.want, func(i, j int) bool {
				return test.want[i].FilePath < test.want[j].FilePath
			})
			for i := range got {
				sort.Strings(got[i].Types)
			}
			for i := range test.want {
				sort.Strings(test.want[i].Types)
			}
			cmpopt := cmpopts.IgnoreFields(licenses.License{}, "Contents")
			if diff := cmp.Diff(test.want, got, cmpopt); diff != "" {
				t.Errorf("mismatch (-want +got):\n%s", diff)
			}
		})
	}
}

func TestGetModuleLicenses(t *testing.T) {
	modulePath := "test.module"
	testModule := sample.LegacyModule(modulePath, "v1.2.3", "", "foo", "bar")
	testModule.LegacyPackages[0].Licenses = []*licenses.Metadata{{Types: []string{"ISC"}, FilePath: "LICENSE"}}
	testModule.LegacyPackages[1].Licenses = []*licenses.Metadata{{Types: []string{"MIT"}, FilePath: "foo/LICENSE"}}
	testModule.LegacyPackages[2].Licenses = []*licenses.Metadata{{Types: []string{"GPL2"}, FilePath: "bar/LICENSE.txt"}}

	defer ResetTestDB(testDB, t)
	ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
	defer cancel()

	testModule.Licenses = nil
	for _, p := range testModule.LegacyPackages {
		testModule.Licenses = append(testModule.Licenses, &licenses.License{
			Metadata: p.Licenses[0],
			Contents: []byte(`Lorem Ipsum`),
		})
	}

	if err := testDB.InsertModule(ctx, testModule); err != nil {
		t.Fatal(err)
	}

	got, err := testDB.getModuleLicenses(ctx, modulePath, testModule.Version)
	if err != nil {
		t.Fatal(err)
	}
	// We only want the top-level license.
	wantLicenses := []*licenses.License{testModule.Licenses[0]}
	if diff := cmp.Diff(wantLicenses, got); diff != "" {
		t.Errorf("testDB.getModuleLicenses(ctx, %q, %q) mismatch (-want +got):\n%s", modulePath, testModule.Version, diff)
	}
}

func TestGetLicensesBypass(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
	defer cancel()
	defer ResetTestDB(testDB, t)

	bypassDB := NewBypassingLicenseCheck(testDB.db)

	// Insert with non-redistributable license contents.
	m := nonRedistributableModule()
	if err := bypassDB.InsertModule(ctx, m); err != nil {
		t.Fatal(err)
	}

	// check reads and the second license in the module and compares it with want.
	check := func(bypass bool, want *licenses.License) {
		t.Helper()
		db := testDB
		if bypass {
			db = bypassDB
		}
		u, err := db.GetUnit(ctx, &internal.UnitMeta{Path: sample.ModulePath, ModulePath: sample.ModulePath, Version: m.Version}, internal.WithLicenses)
		if err != nil {
			t.Fatal(err)
		}
		lics := u.LicenseContents
		if len(lics) != 2 {
			t.Fatal("did not read two licenses")
		}
		if diff := cmp.Diff(want, lics[1]); diff != "" {
			t.Errorf("mismatch (-want, +got):\n%s", diff)
		}
	}

	// Read with license bypass includes non-redistributable license contents.
	check(true, sample.NonRedistributableLicense)

	// Read without license bypass does not include non-redistributable license contents.
	nonRedist := *sample.NonRedistributableLicense
	nonRedist.Contents = nil
	check(false, &nonRedist)
}

func nonRedistributableModule() *internal.Module {
	m := sample.LegacyModule(sample.ModulePath, "v1.2.3", "")
	sample.AddLicense(m, sample.NonRedistributableLicense)
	m.IsRedistributable = false
	m.LegacyPackages[0].IsRedistributable = false
	m.Units[0].IsRedistributable = false
	return m
}
<commit_msg>internal/postgres: replace LegacyPackage in TestGetLicenses<commit_after>// Copyright 2020 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package postgres

import (
	"context"
	"errors"
	"sort"
	"testing"

	"github.com/google/go-cmp/cmp"
	"github.com/google/go-cmp/cmp/cmpopts"
	"golang.org/x/pkgsite/internal"
	"golang.org/x/pkgsite/internal/licenses"
	"golang.org/x/pkgsite/internal/stdlib"
	"golang.org/x/pkgsite/internal/testing/sample"
)

func TestGetLicenses(t *testing.T) {
	testModule := sample.LegacyModule(sample.ModulePath, "v1.2.3", "A/B")
	stdlibModule := sample.LegacyModule(stdlib.ModulePath, "v1.13.0", "cmd/go")
	mit := &licenses.Metadata{Types: []string{"MIT"}, FilePath: "LICENSE"}
	bsd := &licenses.Metadata{Types: []string{"BSD-3-Clause"}, FilePath: "A/B/LICENSE"}

	mitLicense := &licenses.License{Metadata: mit}
	bsdLicense := &licenses.License{Metadata: bsd}
	testModule.Licenses = []*licenses.License{bsdLicense, mitLicense}
	sort.Slice(testModule.Units, func(i, j int) bool {
		return testModule.Units[i].Path < testModule.Units[j].Path
	})

	// github.com/valid/module_name
	testModule.Units[0].Licenses = []*licenses.Metadata{mit}
	// github.com/valid/module_name/A
	testModule.Units[1].Licenses = []*licenses.Metadata{mit}
	// github.com/valid/module_name/A/B
	testModule.Units[2].Licenses = []*licenses.Metadata{mit, bsd}

	defer ResetTestDB(testDB, t)
	ctx, cancel := context.WithTimeout(context.Background(), testTimeout*5)
	defer cancel()
	if err := testDB.InsertModule(ctx, testModule); err != nil {
		t.Fatal(err)
	}
	if err := testDB.InsertModule(ctx, stdlibModule); err != nil {
		t.Fatal(err)
	}
	for _, test := range []struct {
		err                                 error
		name, fullPath, modulePath, version string
		want                                []*licenses.License
	}{
		{
			name:       "module root",
			fullPath:   sample.ModulePath,
			modulePath: sample.ModulePath,
			version:    testModule.Version,
			want:       []*licenses.License{testModule.Licenses[1]},
		},
		{
			name:       "package without license",
			fullPath:   sample.ModulePath + "/A",
			modulePath: sample.ModulePath,
			version:    testModule.Version,
			want:       []*licenses.License{testModule.Licenses[1]},
		},
		{
			name:       "package with additional license",
			fullPath:   sample.ModulePath + "/A/B",
			modulePath: sample.ModulePath,
			version:    testModule.Version,
			want:       testModule.Licenses,
		},
		{
			name:       "stdlib directory",
			fullPath:   "cmd",
			modulePath: stdlib.ModulePath,
			version:    stdlibModule.Version,
			want:       stdlibModule.Licenses,
		},
		{
			name:       "stdlib package",
			fullPath:   "cmd/go",
			modulePath: stdlib.ModulePath,
			version:    stdlibModule.Version,
			want:       stdlibModule.Licenses,
		},
		{
			name:       "stdlib module",
			fullPath:   stdlib.ModulePath,
			modulePath: stdlib.ModulePath,
			version:    stdlibModule.Version,
			want:       stdlibModule.Licenses,
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			u, err := testDB.GetUnit(ctx, &internal.UnitMeta{Path: test.fullPath, ModulePath: test.modulePath, Version: test.version}, internal.WithLicenses)
			if !errors.Is(err, test.err) {
				t.Fatal(err)
			}
			got := u.LicenseContents
			sort.Slice(got, func(i, j int) bool {
				return got[i].FilePath < got[j].FilePath
			})
			sort.Slice(test.want, func(i, j int) bool {
				return test.want[i].FilePath < test.want[j].FilePath
			})
			for i := range got {
				sort.Strings(got[i].Types)
			}
			for i := range test.want {
				sort.Strings(test.want[i].Types)
			}
			cmpopt := cmpopts.IgnoreFields(licenses.License{}, "Contents")
			if diff := cmp.Diff(test.want, got, cmpopt); diff != "" {
				t.Errorf("mismatch (-want +got):\n%s", diff)
			}
		})
	}
}

func TestGetModuleLicenses(t *testing.T) {
	modulePath := "test.module"
	testModule := sample.LegacyModule(modulePath, "v1.2.3", "", "foo", "bar")
	testModule.Packages()[0].Licenses = []*licenses.Metadata{{Types: []string{"ISC"}, FilePath: "LICENSE"}}
	testModule.Packages()[1].Licenses = []*licenses.Metadata{{Types: []string{"MIT"}, FilePath: "foo/LICENSE"}}
	testModule.Packages()[2].Licenses = []*licenses.Metadata{{Types: []string{"GPL2"}, FilePath: "bar/LICENSE.txt"}}

	defer ResetTestDB(testDB, t)
	ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
	defer cancel()

	testModule.Licenses = nil
	for _, p := range testModule.Packages() {
		testModule.Licenses = append(testModule.Licenses, &licenses.License{
			Metadata: p.Licenses[0],
			Contents: []byte(`Lorem Ipsum`),
		})
	}

	if err := testDB.InsertModule(ctx, testModule); err != nil {
		t.Fatal(err)
	}

	got, err := testDB.getModuleLicenses(ctx, modulePath, testModule.Version)
	if err != nil {
		t.Fatal(err)
	}
	// We only want the top-level license.
	wantLicenses := []*licenses.License{testModule.Licenses[0]}
	if diff := cmp.Diff(wantLicenses, got); diff != "" {
		t.Errorf("testDB.getModuleLicenses(ctx, %q, %q) mismatch (-want +got):\n%s", modulePath, testModule.Version, diff)
	}
}

func TestGetLicensesBypass(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
	defer cancel()
	defer ResetTestDB(testDB, t)

	bypassDB := NewBypassingLicenseCheck(testDB.db)

	// Insert with non-redistributable license contents.
	m := nonRedistributableModule()
	if err := bypassDB.InsertModule(ctx, m); err != nil {
		t.Fatal(err)
	}

	// check reads and the second license in the module and compares it with want.
	check := func(bypass bool, want *licenses.License) {
		t.Helper()
		db := testDB
		if bypass {
			db = bypassDB
		}
		u, err := db.GetUnit(ctx, &internal.UnitMeta{Path: sample.ModulePath, ModulePath: sample.ModulePath, Version: m.Version}, internal.WithLicenses)
		if err != nil {
			t.Fatal(err)
		}
		lics := u.LicenseContents
		if len(lics) != 2 {
			t.Fatal("did not read two licenses")
		}
		if diff := cmp.Diff(want, lics[1]); diff != "" {
			t.Errorf("mismatch (-want, +got):\n%s", diff)
		}
	}

	// Read with license bypass includes non-redistributable license contents.
	check(true, sample.NonRedistributableLicense)

	// Read without license bypass does not include non-redistributable license contents.
	nonRedist := *sample.NonRedistributableLicense
	nonRedist.Contents = nil
	check(false, &nonRedist)
}

func nonRedistributableModule() *internal.Module {
	m := sample.LegacyModule(sample.ModulePath, "v1.2.3", "")
	sample.AddLicense(m, sample.NonRedistributableLicense)
	m.IsRedistributable = false
	m.Packages()[0].IsRedistributable = false
	m.Units[0].IsRedistributable = false
	return m
}
<|endoftext|>