<commit_before>package lars

import (
	"strconv"
	"strings"

	"github.com/gorilla/websocket"
)

// IRouteGroup interface for router group
type IRouteGroup interface {
	IRoutes
	Group(prefix string, middleware ...Handler) IRouteGroup
}

// IRoutes interface for routes
type IRoutes interface {
	Use(...Handler)
	Any(string, ...Handler)
	Get(string, ...Handler)
	Post(string, ...Handler)
	Delete(string, ...Handler)
	Patch(string, ...Handler)
	Put(string, ...Handler)
	Options(string, ...Handler)
	Head(string, ...Handler)
	Connect(string, ...Handler)
	Trace(string, ...Handler)
	WebSocket(websocket.Upgrader, string, Handler)
}

// routeGroup struct containing all fields and methods for use.
type routeGroup struct {
	prefix     string
	middleware HandlersChain
	lars       *LARS
}

var _ IRouteGroup = &routeGroup{}

func (g *routeGroup) handle(method string, path string, handlers []Handler) {

	if len(handlers) == 0 {
		panic("No handler mapped to path:" + path)
	}

	if i := strings.Index(path, "//"); i != -1 {
		panic("Bad path '" + path + "' contains duplicate // at index:" + strconv.Itoa(i))
	}

	chain := make(HandlersChain, len(handlers))
	name := ""

	for i, h := range handlers {

		if i == len(handlers)-1 {
			chain[i], name = g.lars.wrapHandlerWithName(h)
		} else {
			chain[i] = g.lars.wrapHandler(h)
		}
	}

	tree := g.lars.trees[method]
	if tree == nil {
		tree = new(node)
		g.lars.trees[method] = tree
	}

	combined := make(HandlersChain, len(g.middleware)+len(chain))
	copy(combined, g.middleware)
	copy(combined[len(g.middleware):], chain)

	pCount := tree.add(g.prefix+path, name, combined)
	pCount++

	if pCount > g.lars.mostParams {
		g.lars.mostParams = pCount
	}
}

// Use adds a middleware handler to the group middleware chain.
func (g *routeGroup) Use(m ...Handler) {
	for _, h := range m {
		g.middleware = append(g.middleware, g.lars.wrapHandler(h))
	}
}

// Connect adds a CONNECT route & handler to the router.
func (g *routeGroup) Connect(path string, h ...Handler) {
	g.handle(CONNECT, path, h)
}

// Delete adds a DELETE route & handler to the router.
func (g *routeGroup) Delete(path string, h ...Handler) {
	g.handle(DELETE, path, h)
}

// Get adds a GET route & handler to the router.
func (g *routeGroup) Get(path string, h ...Handler) {
	g.handle(GET, path, h)
}

// Head adds a HEAD route & handler to the router.
func (g *routeGroup) Head(path string, h ...Handler) {
	g.handle(HEAD, path, h)
}

// Options adds an OPTIONS route & handler to the router.
func (g *routeGroup) Options(path string, h ...Handler) {
	g.handle(OPTIONS, path, h)
}

// Patch adds a PATCH route & handler to the router.
func (g *routeGroup) Patch(path string, h ...Handler) {
	g.handle(PATCH, path, h)
}

// Post adds a POST route & handler to the router.
func (g *routeGroup) Post(path string, h ...Handler) {
	g.handle(POST, path, h)
}

// Put adds a PUT route & handler to the router.
func (g *routeGroup) Put(path string, h ...Handler) {
	g.handle(PUT, path, h)
}

// Trace adds a TRACE route & handler to the router.
func (g *routeGroup) Trace(path string, h ...Handler) {
	g.handle(TRACE, path, h)
}

// Handle allows for any method to be registered with the given
// route & handler. Allows for non standard methods to be used
// like CalDavs PROPFIND and so forth.
func (g *routeGroup) Handle(method string, path string, h ...Handler) {
	g.handle(method, path, h)
}

// Any adds a route & handler to the router for all HTTP methods.
func (g *routeGroup) Any(path string, h ...Handler) {
	g.Connect(path, h...)
	g.Delete(path, h...)
	g.Get(path, h...)
	g.Head(path, h...)
	g.Options(path, h...)
	g.Patch(path, h...)
	g.Post(path, h...)
	g.Put(path, h...)
	g.Trace(path, h...)
}

// Match adds a route & handler to the router for multiple HTTP methods provided.
func (g *routeGroup) Match(methods []string, path string, h ...Handler) {
	for _, m := range methods {
		g.handle(m, path, h)
	}
}

// WebSocket adds a websocket route
func (g *routeGroup) WebSocket(upgrader websocket.Upgrader, path string, h Handler) {

	handler := g.lars.wrapHandler(h)
	g.Get(path, func(c Context) {

		ctx := c.BaseContext()
		var err error

		ctx.websocket, err = upgrader.Upgrade(ctx.response, ctx.request, nil)
		if err != nil {
			return
		}

		defer ctx.websocket.Close()
		c.Next()
	}, handler)
}

// Group creates a new sub router with prefix. It inherits all properties from
// the parent. Passing middleware overrides parent middleware but still keeps
// the root level middleware intact.
func (g *routeGroup) Group(prefix string, middleware ...Handler) IRouteGroup {

	rg := &routeGroup{
		prefix: g.prefix + prefix,
		lars:   g.lars,
	}

	if len(middleware) == 0 {
		rg.middleware = make(HandlersChain, len(g.middleware)+len(middleware))
		copy(rg.middleware, g.middleware)

		return rg
	}

	if middleware[0] == nil {
		rg.middleware = make(HandlersChain, 0)
		return rg
	}

	rg.middleware = make(HandlersChain, len(middleware))
	copy(rg.middleware, g.lars.middleware)
	rg.Use(middleware...)

	return rg
}
<commit_msg>Fix middleware copying when creating new group<commit_after>package lars

import (
	"strconv"
	"strings"

	"github.com/gorilla/websocket"
)

// IRouteGroup interface for router group
type IRouteGroup interface {
	IRoutes
	Group(prefix string, middleware ...Handler) IRouteGroup
}

// IRoutes interface for routes
type IRoutes interface {
	Use(...Handler)
	Any(string, ...Handler)
	Get(string, ...Handler)
	Post(string, ...Handler)
	Delete(string, ...Handler)
	Patch(string, ...Handler)
	Put(string, ...Handler)
	Options(string, ...Handler)
	Head(string, ...Handler)
	Connect(string, ...Handler)
	Trace(string, ...Handler)
	WebSocket(websocket.Upgrader, string, Handler)
}

// routeGroup struct containing all fields and methods for use.
type routeGroup struct {
	prefix     string
	middleware HandlersChain
	lars       *LARS
}

var _ IRouteGroup = &routeGroup{}

func (g *routeGroup) handle(method string, path string, handlers []Handler) {

	if len(handlers) == 0 {
		panic("No handler mapped to path:" + path)
	}

	if i := strings.Index(path, "//"); i != -1 {
		panic("Bad path '" + path + "' contains duplicate // at index:" + strconv.Itoa(i))
	}

	chain := make(HandlersChain, len(handlers))
	name := ""

	for i, h := range handlers {

		if i == len(handlers)-1 {
			chain[i], name = g.lars.wrapHandlerWithName(h)
		} else {
			chain[i] = g.lars.wrapHandler(h)
		}
	}

	tree := g.lars.trees[method]
	if tree == nil {
		tree = new(node)
		g.lars.trees[method] = tree
	}

	combined := make(HandlersChain, len(g.middleware)+len(chain))
	copy(combined, g.middleware)
	copy(combined[len(g.middleware):], chain)

	pCount := tree.add(g.prefix+path, name, combined)
	pCount++

	if pCount > g.lars.mostParams {
		g.lars.mostParams = pCount
	}
}

// Use adds a middleware handler to the group middleware chain.
func (g *routeGroup) Use(m ...Handler) {
	for _, h := range m {
		g.middleware = append(g.middleware, g.lars.wrapHandler(h))
	}
}

// Connect adds a CONNECT route & handler to the router.
func (g *routeGroup) Connect(path string, h ...Handler) {
	g.handle(CONNECT, path, h)
}

// Delete adds a DELETE route & handler to the router.
func (g *routeGroup) Delete(path string, h ...Handler) {
	g.handle(DELETE, path, h)
}

// Get adds a GET route & handler to the router.
func (g *routeGroup) Get(path string, h ...Handler) {
	g.handle(GET, path, h)
}

// Head adds a HEAD route & handler to the router.
func (g *routeGroup) Head(path string, h ...Handler) {
	g.handle(HEAD, path, h)
}

// Options adds an OPTIONS route & handler to the router.
func (g *routeGroup) Options(path string, h ...Handler) {
	g.handle(OPTIONS, path, h)
}

// Patch adds a PATCH route & handler to the router.
func (g *routeGroup) Patch(path string, h ...Handler) {
	g.handle(PATCH, path, h)
}

// Post adds a POST route & handler to the router.
func (g *routeGroup) Post(path string, h ...Handler) {
	g.handle(POST, path, h)
}

// Put adds a PUT route & handler to the router.
func (g *routeGroup) Put(path string, h ...Handler) {
	g.handle(PUT, path, h)
}

// Trace adds a TRACE route & handler to the router.
func (g *routeGroup) Trace(path string, h ...Handler) {
	g.handle(TRACE, path, h)
}

// Handle allows for any method to be registered with the given
// route & handler. Allows for non standard methods to be used
// like CalDavs PROPFIND and so forth.
func (g *routeGroup) Handle(method string, path string, h ...Handler) {
	g.handle(method, path, h)
}

// Any adds a route & handler to the router for all HTTP methods.
func (g *routeGroup) Any(path string, h ...Handler) {
	g.Connect(path, h...)
	g.Delete(path, h...)
	g.Get(path, h...)
	g.Head(path, h...)
	g.Options(path, h...)
	g.Patch(path, h...)
	g.Post(path, h...)
	g.Put(path, h...)
	g.Trace(path, h...)
}

// Match adds a route & handler to the router for multiple HTTP methods provided.
func (g *routeGroup) Match(methods []string, path string, h ...Handler) {
	for _, m := range methods {
		g.handle(m, path, h)
	}
}

// WebSocket adds a websocket route
func (g *routeGroup) WebSocket(upgrader websocket.Upgrader, path string, h Handler) {

	handler := g.lars.wrapHandler(h)
	g.Get(path, func(c Context) {

		ctx := c.BaseContext()
		var err error

		ctx.websocket, err = upgrader.Upgrade(ctx.response, ctx.request, nil)
		if err != nil {
			return
		}

		defer ctx.websocket.Close()
		c.Next()
	}, handler)
}

// Group creates a new sub router with prefix. It inherits all properties from
// the parent. Passing middleware overrides parent middleware but still keeps
// the root level middleware intact.
func (g *routeGroup) Group(prefix string, middleware ...Handler) IRouteGroup {

	rg := &routeGroup{
		prefix: g.prefix + prefix,
		lars:   g.lars,
	}

	if len(middleware) == 0 {
		rg.middleware = make(HandlersChain, len(g.middleware)+len(middleware))
		copy(rg.middleware, g.middleware)

		return rg
	}

	if middleware[0] == nil {
		rg.middleware = make(HandlersChain, 0)
		return rg
	}

	rg.middleware = make(HandlersChain, len(g.lars.middleware))
	copy(rg.middleware, g.lars.middleware)
	rg.Use(middleware...)

	return rg
}
<|endoftext|>


<commit_before>package conn

import (
	"fmt"
	"io"
	"net"

	"QmSir6qPL1tjuxd8LkR8VZq6v625ExAUVs2eCLeqQuaPGU/goprocess"
	goprocessctx "QmSir6qPL1tjuxd8LkR8VZq6v625ExAUVs2eCLeqQuaPGU/goprocess/context"
	context "QmacZi9WygGK7Me8mH53pypyscHzU386aUZXpr28GZgUct/context"
	ma "QmbWxL1aXQhBjc1XGjGF1f2KGBMCBYSuT2ThA8YXnXJK83/go-multiaddr"
	tec "github.com/jbenet/go-temp-err-catcher"

	ic "github.com/ipfs/go-libp2p/p2p/crypto"
	filter "github.com/ipfs/go-libp2p/p2p/net/filter"
	transport "github.com/ipfs/go-libp2p/p2p/net/transport"
	peer "github.com/ipfs/go-libp2p/p2p/peer"
)

// ConnWrapper is any function that wraps a raw multiaddr connection
type ConnWrapper func(transport.Conn) transport.Conn

// listener is an object that can accept connections. It implements Listener
type listener struct {
	transport.Listener

	local peer.ID    // LocalPeer is the identity of the local Peer
	privk ic.PrivKey // private key to use to initialize secure conns

	filters *filter.Filters

	wrapper ConnWrapper

	proc goprocess.Process
}

func (l *listener) teardown() error {
	defer log.Debugf("listener closed: %s %s", l.local, l.Multiaddr())
	return l.Listener.Close()
}

func (l *listener) Close() error {
	log.Debugf("listener closing: %s %s", l.local, l.Multiaddr())
	return l.proc.Close()
}

func (l *listener) String() string {
	return fmt.Sprintf("<Listener %s %s>", l.local, l.Multiaddr())
}

func (l *listener) SetAddrFilters(fs *filter.Filters) {
	l.filters = fs
}

// Accept waits for and returns the next connection to the listener.
// Note that unfortunately this
func (l *listener) Accept() (net.Conn, error) {

	// listeners dont have contexts. given changes dont make sense here anymore
	// note that the parent of listener will Close, which will interrupt all io.
	// Contexts and io don't mix.
	ctx := context.Background()

	var catcher tec.TempErrCatcher

	catcher.IsTemp = func(e error) bool {
		// ignore connection breakages up to this point. but log them
		if e == io.EOF {
			log.Debugf("listener ignoring conn with EOF: %s", e)
			return true
		}

		te, ok := e.(tec.Temporary)
		if ok {
			log.Debugf("listener ignoring conn with temporary err: %s", e)
			return te.Temporary()
		}
		return false
	}

	for {
		maconn, err := l.Listener.Accept()
		if err != nil {
			if catcher.IsTemporary(err) {
				continue
			}
			return nil, err
		}

		log.Debugf("listener %s got connection: %s <---> %s", l, maconn.LocalMultiaddr(), maconn.RemoteMultiaddr())

		if l.filters != nil && l.filters.AddrBlocked(maconn.RemoteMultiaddr()) {
			log.Debugf("blocked connection from %s", maconn.RemoteMultiaddr())
			maconn.Close()
			continue
		}
		// If we have a wrapper func, wrap this conn
		if l.wrapper != nil {
			maconn = l.wrapper(maconn)
		}

		c, err := newSingleConn(ctx, l.local, "", maconn)
		if err != nil {
			if catcher.IsTemporary(err) {
				continue
			}
			return nil, err
		}

		if l.privk == nil || EncryptConnections == false {
			log.Warning("listener %s listening INSECURELY!", l)
			return c, nil
		}
		sc, err := newSecureConn(ctx, l.privk, c)
		if err != nil {
			log.Infof("ignoring conn we failed to secure: %s %s", err, c)
			continue
		}
		return sc, nil
	}
}

func (l *listener) Addr() net.Addr {
	return l.Listener.Addr()
}

// Multiaddr is the identity of the local Peer.
// If there is an error converting from net.Addr to ma.Multiaddr,
// the return value will be nil.
func (l *listener) Multiaddr() ma.Multiaddr {
	return l.Listener.Multiaddr()
}

// LocalPeer is the identity of the local Peer.
func (l *listener) LocalPeer() peer.ID {
	return l.local
}

func (l *listener) Loggable() map[string]interface{} {
	return map[string]interface{}{
		"listener": map[string]interface{}{
			"peer":    l.LocalPeer(),
			"address": l.Multiaddr(),
			"secure":  (l.privk != nil),
		},
	}
}

func WrapTransportListener(ctx context.Context, ml transport.Listener, local peer.ID, sk ic.PrivKey) (Listener, error) {
	l := &listener{
		Listener: ml,
		local:    local,
		privk:    sk,
	}
	l.proc = goprocessctx.WithContextAndTeardown(ctx, l.teardown)

	log.Debugf("Conn Listener on %s", l.Multiaddr())
	log.Event(ctx, "swarmListen", l)
	return l, nil
}

type ListenerConnWrapper interface {
	SetConnWrapper(ConnWrapper)
}

// SetConnWrapper assigns a maconn ConnWrapper to wrap all incoming
// connections with. MUST be set _before_ calling `Accept()`
func (l *listener) SetConnWrapper(cw ConnWrapper) {
	l.wrapper = cw
}
<commit_msg>Vendor in go-peerstream<commit_after>package conn

import (
	"fmt"
	"io"
	"net"

	"QmSir6qPL1tjuxd8LkR8VZq6v625ExAUVs2eCLeqQuaPGU/goprocess"
	goprocessctx "QmSir6qPL1tjuxd8LkR8VZq6v625ExAUVs2eCLeqQuaPGU/goprocess/context"
	tec "QmWtLNgjHvFnRHcHUheAMGx4sLYYYGSacNA3eG52ywy2UQ/go-temp-err-catcher"
	context "QmacZi9WygGK7Me8mH53pypyscHzU386aUZXpr28GZgUct/context"
	ma "QmbWxL1aXQhBjc1XGjGF1f2KGBMCBYSuT2ThA8YXnXJK83/go-multiaddr"

	ic "github.com/ipfs/go-libp2p/p2p/crypto"
	filter "github.com/ipfs/go-libp2p/p2p/net/filter"
	transport "github.com/ipfs/go-libp2p/p2p/net/transport"
	peer "github.com/ipfs/go-libp2p/p2p/peer"
)

// ConnWrapper is any function that wraps a raw multiaddr connection
type ConnWrapper func(transport.Conn) transport.Conn

// listener is an object that can accept connections. It implements Listener
type listener struct {
	transport.Listener

	local peer.ID    // LocalPeer is the identity of the local Peer
	privk ic.PrivKey // private key to use to initialize secure conns

	filters *filter.Filters

	wrapper ConnWrapper

	proc goprocess.Process
}

func (l *listener) teardown() error {
	defer log.Debugf("listener closed: %s %s", l.local, l.Multiaddr())
	return l.Listener.Close()
}

func (l *listener) Close() error {
	log.Debugf("listener closing: %s %s", l.local, l.Multiaddr())
	return l.proc.Close()
}

func (l *listener) String() string {
	return fmt.Sprintf("<Listener %s %s>", l.local, l.Multiaddr())
}

func (l *listener) SetAddrFilters(fs *filter.Filters) {
	l.filters = fs
}

// Accept waits for and returns the next connection to the listener.
// Note that unfortunately this
func (l *listener) Accept() (net.Conn, error) {

	// listeners dont have contexts. given changes dont make sense here anymore
	// note that the parent of listener will Close, which will interrupt all io.
	// Contexts and io don't mix.
	ctx := context.Background()

	var catcher tec.TempErrCatcher

	catcher.IsTemp = func(e error) bool {
		// ignore connection breakages up to this point. but log them
		if e == io.EOF {
			log.Debugf("listener ignoring conn with EOF: %s", e)
			return true
		}

		te, ok := e.(tec.Temporary)
		if ok {
			log.Debugf("listener ignoring conn with temporary err: %s", e)
			return te.Temporary()
		}
		return false
	}

	for {
		maconn, err := l.Listener.Accept()
		if err != nil {
			if catcher.IsTemporary(err) {
				continue
			}
			return nil, err
		}

		log.Debugf("listener %s got connection: %s <---> %s", l, maconn.LocalMultiaddr(), maconn.RemoteMultiaddr())

		if l.filters != nil && l.filters.AddrBlocked(maconn.RemoteMultiaddr()) {
			log.Debugf("blocked connection from %s", maconn.RemoteMultiaddr())
			maconn.Close()
			continue
		}
		// If we have a wrapper func, wrap this conn
		if l.wrapper != nil {
			maconn = l.wrapper(maconn)
		}

		c, err := newSingleConn(ctx, l.local, "", maconn)
		if err != nil {
			if catcher.IsTemporary(err) {
				continue
			}
			return nil, err
		}

		if l.privk == nil || EncryptConnections == false {
			log.Warning("listener %s listening INSECURELY!", l)
			return c, nil
		}
		sc, err := newSecureConn(ctx, l.privk, c)
		if err != nil {
			log.Infof("ignoring conn we failed to secure: %s %s", err, c)
			continue
		}
		return sc, nil
	}
}

func (l *listener) Addr() net.Addr {
	return l.Listener.Addr()
}

// Multiaddr is the identity of the local Peer.
// If there is an error converting from net.Addr to ma.Multiaddr,
// the return value will be nil.
func (l *listener) Multiaddr() ma.Multiaddr {
	return l.Listener.Multiaddr()
}

// LocalPeer is the identity of the local Peer.
func (l *listener) LocalPeer() peer.ID {
	return l.local
}

func (l *listener) Loggable() map[string]interface{} {
	return map[string]interface{}{
		"listener": map[string]interface{}{
			"peer":    l.LocalPeer(),
			"address": l.Multiaddr(),
			"secure":  (l.privk != nil),
		},
	}
}

func WrapTransportListener(ctx context.Context, ml transport.Listener, local peer.ID, sk ic.PrivKey) (Listener, error) {
	l := &listener{
		Listener: ml,
		local:    local,
		privk:    sk,
	}
	l.proc = goprocessctx.WithContextAndTeardown(ctx, l.teardown)

	log.Debugf("Conn Listener on %s", l.Multiaddr())
	log.Event(ctx, "swarmListen", l)
	return l, nil
}

type ListenerConnWrapper interface {
	SetConnWrapper(ConnWrapper)
}

// SetConnWrapper assigns a maconn ConnWrapper to wrap all incoming
// connections with. MUST be set _before_ calling `Accept()`
func (l *listener) SetConnWrapper(cw ConnWrapper) {
	l.wrapper = cw
}
<|endoftext|>


<commit_before>// Copyright 2016 CoreOS, 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,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build cgo

package util

// #include <stdlib.h>
// #include <sys/types.h>
// #include <unistd.h>
//
// int
// my_sd_pid_get_owner_uid(void *f, pid_t pid, uid_t *uid)
// {
//   int (*sd_pid_get_owner_uid)(pid_t, uid_t *);
//
//   sd_pid_get_owner_uid = (int (*)(pid_t, uid_t *))f;
//   return sd_pid_get_owner_uid(pid, uid);
// }
//
// int
// my_sd_pid_get_unit(void *f, pid_t pid, char **unit)
// {
//   int (*sd_pid_get_unit)(pid_t, char **);
//
//   sd_pid_get_unit = (int (*)(pid_t, char **))f;
//   return sd_pid_get_unit(pid, unit);
// }
//
// int
// my_sd_pid_get_slice(void *f, pid_t pid, char **slice)
// {
//   int (*sd_pid_get_slice)(pid_t, char **);
//
//   sd_pid_get_slice = (int (*)(pid_t, char **))f;
//   return sd_pid_get_slice(pid, slice);
// }
//
// int
// am_session_leader()
// {
//   return (getsid(0) == getpid());
// }
import "C"
import (
	"fmt"
	"syscall"
	"unsafe"

	"github.com/coreos/pkg/dlopen"
)

var libsystemdNames = []string{
	// systemd < 209
	"libsystemd-login.so.0",
	"libsystemd-login.so",

	// systemd >= 209 merged libsystemd-login into libsystemd proper
	"libsystemd.so.0",
	"libsystemd.so",
}

func getRunningSlice() (slice string, err error) {
	var h *dlopen.LibHandle
	h, err = dlopen.GetHandle(libsystemdNames)
	if err != nil {
		return
	}
	defer func() {
		if err1 := h.Close(); err1 != nil {
			err = err1
		}
	}()

	sd_pid_get_slice, err := h.GetSymbolPointer("sd_pid_get_slice")
	if err != nil {
		return
	}

	var s string
	sl := C.CString(s)
	defer C.free(unsafe.Pointer(sl))

	ret := C.my_sd_pid_get_slice(sd_pid_get_slice, 0, &sl)
	if ret < 0 {
		err = fmt.Errorf("error calling sd_pid_get_slice: %v", syscall.Errno(-ret))
		return
	}

	return C.GoString(sl), nil
}

func runningFromSystemService() (ret bool, err error) {
	var h *dlopen.LibHandle
	h, err = dlopen.GetHandle(libsystemdNames)
	if err != nil {
		return
	}
	defer func() {
		if err1 := h.Close(); err1 != nil {
			err = err1
		}
	}()

	sd_pid_get_owner_uid, err := h.GetSymbolPointer("sd_pid_get_owner_uid")
	if err != nil {
		return
	}

	var uid C.uid_t
	errno := C.my_sd_pid_get_owner_uid(sd_pid_get_owner_uid, 0, &uid)
	serrno := syscall.Errno(-errno)
	// when we're running from a unit file, sd_pid_get_owner_uid returns
	// ENOENT (systemd <220) or ENXIO (systemd >=220)
	switch {
	case errno >= 0:
		ret = false
	case serrno == syscall.ENOENT, serrno == syscall.ENXIO:
		// Since the implementation of sessions in systemd relies on
		// the `pam_systemd` module, using the sd_pid_get_owner_uid
		// heuristic alone can result in false positives if that module
		// (or PAM itself) is not present or properly configured on the
		// system. As such, we also check if we're the session leader,
		// which should be the case if we're invoked from a unit file,
		// but not if e.g. we're invoked from the command line from a
		// user's login session
		ret = C.am_session_leader() == 1
	default:
		err = fmt.Errorf("error calling sd_pid_get_owner_uid: %v", syscall.Errno(-errno))
	}
	return
}

func currentUnitName() (unit string, err error) {
	var h *dlopen.LibHandle
	h, err = dlopen.GetHandle(libsystemdNames)
	if err != nil {
		return
	}
	defer func() {
		if err1 := h.Close(); err1 != nil {
			err = err1
		}
	}()

	sd_pid_get_unit, err := h.GetSymbolPointer("sd_pid_get_unit")
	if err != nil {
		return
	}

	var s string
	u := C.CString(s)
	defer C.free(unsafe.Pointer(u))

	ret := C.my_sd_pid_get_unit(sd_pid_get_unit, 0, &u)
	if ret < 0 {
		err = fmt.Errorf("error calling sd_pid_get_unit: %v", syscall.Errno(-ret))
		return
	}

	unit = C.GoString(u)
	return
}
<commit_msg>util: fix runningFromSystemService<commit_after>// Copyright 2016 CoreOS, 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,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build cgo

package util

// #include <stdlib.h>
// #include <sys/types.h>
// #include <unistd.h>
//
// int
// my_sd_pid_get_owner_uid(void *f, pid_t pid, uid_t *uid)
// {
//   int (*sd_pid_get_owner_uid)(pid_t, uid_t *);
//
//   sd_pid_get_owner_uid = (int (*)(pid_t, uid_t *))f;
//   return sd_pid_get_owner_uid(pid, uid);
// }
//
// int
// my_sd_pid_get_unit(void *f, pid_t pid, char **unit)
// {
//   int (*sd_pid_get_unit)(pid_t, char **);
//
//   sd_pid_get_unit = (int (*)(pid_t, char **))f;
//   return sd_pid_get_unit(pid, unit);
// }
//
// int
// my_sd_pid_get_slice(void *f, pid_t pid, char **slice)
// {
//   int (*sd_pid_get_slice)(pid_t, char **);
//
//   sd_pid_get_slice = (int (*)(pid_t, char **))f;
//   return sd_pid_get_slice(pid, slice);
// }
//
// int
// am_session_leader()
// {
//   return (getsid(0) == getpid());
// }
import "C"
import (
	"fmt"
	"syscall"
	"unsafe"

	"github.com/coreos/pkg/dlopen"
)

var libsystemdNames = []string{
	// systemd < 209
	"libsystemd-login.so.0",
	"libsystemd-login.so",

	// systemd >= 209 merged libsystemd-login into libsystemd proper
	"libsystemd.so.0",
	"libsystemd.so",
}

func getRunningSlice() (slice string, err error) {
	var h *dlopen.LibHandle
	h, err = dlopen.GetHandle(libsystemdNames)
	if err != nil {
		return
	}
	defer func() {
		if err1 := h.Close(); err1 != nil {
			err = err1
		}
	}()

	sd_pid_get_slice, err := h.GetSymbolPointer("sd_pid_get_slice")
	if err != nil {
		return
	}

	var s string
	sl := C.CString(s)
	defer C.free(unsafe.Pointer(sl))

	ret := C.my_sd_pid_get_slice(sd_pid_get_slice, 0, &sl)
	if ret < 0 {
		err = fmt.Errorf("error calling sd_pid_get_slice: %v", syscall.Errno(-ret))
		return
	}

	return C.GoString(sl), nil
}

func runningFromSystemService() (ret bool, err error) {
	var h *dlopen.LibHandle
	h, err = dlopen.GetHandle(libsystemdNames)
	if err != nil {
		return
	}
	defer func() {
		if err1 := h.Close(); err1 != nil {
			err = err1
		}
	}()

	sd_pid_get_owner_uid, err := h.GetSymbolPointer("sd_pid_get_owner_uid")
	if err != nil {
		return
	}

	var uid C.uid_t
	errno := C.my_sd_pid_get_owner_uid(sd_pid_get_owner_uid, 0, &uid)
	serrno := syscall.Errno(-errno)
	// when we're running from a unit file, sd_pid_get_owner_uid returns
	// ENOENT (systemd <220), ENXIO (systemd 220-223), or ENODATA
	// (systemd >=234)
	switch {
	case errno >= 0:
		ret = false
	case serrno == syscall.ENOENT, serrno == syscall.ENXIO, serrno == syscall.ENODATA:
		// Since the implementation of sessions in systemd relies on
		// the `pam_systemd` module, using the sd_pid_get_owner_uid
		// heuristic alone can result in false positives if that module
		// (or PAM itself) is not present or properly configured on the
		// system. As such, we also check if we're the session leader,
		// which should be the case if we're invoked from a unit file,
		// but not if e.g. we're invoked from the command line from a
		// user's login session
		ret = C.am_session_leader() == 1
	default:
		err = fmt.Errorf("error calling sd_pid_get_owner_uid: %v", syscall.Errno(-errno))
	}
	return
}

func currentUnitName() (unit string, err error) {
	var h *dlopen.LibHandle
	h, err = dlopen.GetHandle(libsystemdNames)
	if err != nil {
		return
	}
	defer func() {
		if err1 := h.Close(); err1 != nil {
			err = err1
		}
	}()

	sd_pid_get_unit, err := h.GetSymbolPointer("sd_pid_get_unit")
	if err != nil {
		return
	}

	var s string
	u := C.CString(s)
	defer C.free(unsafe.Pointer(u))

	ret := C.my_sd_pid_get_unit(sd_pid_get_unit, 0, &u)
	if ret < 0 {
		err = fmt.Errorf("error calling sd_pid_get_unit: %v", syscall.Errno(-ret))
		return
	}

	unit = C.GoString(u)
	return
}
<|endoftext|>


<commit_before>package session

import (
	"fmt"

	"github.com/asaskevich/govalidator"
	"github.com/ellcrys/consul/api"
	"github.com/ellcrys/util"
	"github.com/pkg/errors"
)

var SessionTTL = "900s"

// ConsulRegistry implements a session registry based on
// consul. It satisfies SessionRegistry interface.
type ConsulRegistry struct {
	client *api.Client
}

// NewConsulRegistry creates a new consul registry.
// Connection to consul is attempted, error is returned on failure
// to connect.
func NewConsulRegistry() (r *ConsulRegistry, err error) {
	r = new(ConsulRegistry)
	cfg := api.DefaultConfig()
	cfg.Address = util.Env("CONSUL_ADDR", "127.0.0.1:8500")
	r.client, err = api.NewClient(cfg)
	if err != nil {
		err = errors.Wrap(err, "failed to create client")
		return
	}
	if _, err = r.client.Status().Leader(); err != nil {
		err = errors.Wrap(err, "failed to connect to consul")
	}
	return
}

// makeKey creates a key to be used in consul KV store
func makeKey(sid string) string {
	return "db_session_" + sid
}

// Add registers a new session
func (r *ConsulRegistry) Add(item RegItem) error {

	if !govalidator.IsUUIDv4(item.SID) {
		return fmt.Errorf("session id must be a UUID")
	}

	session := r.client.Session()
	sessionID, _, err := session.Create(&api.SessionEntry{
		ID:       item.SID,
		TTL:      SessionTTL,
		Behavior: api.SessionBehaviorDelete,
	}, nil)
	if err != nil {
		return errors.Wrap(err, "failed to create consul session")
	}

	acquired, _, err := r.client.KV().Acquire(&api.KVPair{
		Key:     makeKey(item.SID),
		Value:   item.ToJSON(),
		Session: sessionID,
	}, nil)
	if err != nil {
		return errors.Wrap(err, "failed to create consul key for db session")
	}

	if !acquired {
		return fmt.Errorf("failed to acquire lock on session")
	}

	return err
}

// Get gets a session. Returns ErrNotFound if not found
func (r *ConsulRegistry) Get(sid string) (*RegItem, error) {

	kv, _, err := r.client.KV().Get(makeKey(sid), nil)
	if err != nil {
		return nil, errors.Wrap(err, "failed to get session")
	}
	if kv == nil {
		return nil, ErrNotFound
	}
	var item RegItem
	if err = util.FromJSON(kv.Value, &item); err != nil {
		return nil, errors.Wrap(err, "failed to parse registry item data")
	}
	return &item, nil
}

// Del deletes a session. Returns nil if the session was
// deleted successfully or it doesn't exists
func (r *ConsulRegistry) Del(sid string) (err error) {

	if !govalidator.IsUUIDv4(sid) {
		return fmt.Errorf("session id must be a UUID")
	}

	_, err = r.client.Session().Destroy(sid, nil)
	if err != nil {
		return errors.Wrap(err, "failed to delete registry item")
	}

	_, err = r.client.KV().Delete(makeKey(sid), nil)
	if err != nil {
		return errors.Wrap(err, "failed to delete session key")
	}

	return
}
<commit_msg>use official consul API<commit_after>package session

import (
	"fmt"

	"github.com/asaskevich/govalidator"
	"github.com/ellcrys/util"
	"github.com/hashicorp/consul/api"
	"github.com/pkg/errors"
)

var SessionTTL = "900s"

// ConsulRegistry implements a session registry based on
// consul. It satisfies SessionRegistry interface.
type ConsulRegistry struct {
	client *api.Client
}

// NewConsulRegistry creates a new consul registry.
// Connection to consul is attempted, error is returned on failure
// to connect.
func NewConsulRegistry() (r *ConsulRegistry, err error) {
	r = new(ConsulRegistry)
	cfg := api.DefaultConfig()
	cfg.Address = util.Env("CONSUL_ADDR", "127.0.0.1:8500")
	r.client, err = api.NewClient(cfg)
	if err != nil {
		err = errors.Wrap(err, "failed to create client")
		return
	}
	if _, err = r.client.Status().Leader(); err != nil {
		err = errors.Wrap(err, "failed to connect to consul")
	}
	return
}

// makeKey creates a key to be used in consul KV store
func makeKey(sid string) string {
	return "db_session_" + sid
}

// Add registers a new session
func (r *ConsulRegistry) Add(item RegItem) error {

	if !govalidator.IsUUIDv4(item.SID) {
		return fmt.Errorf("session id must be a UUID")
	}

	session := r.client.Session()
	sessionID, _, err := session.Create(&api.SessionEntry{
		ID:       item.SID,
		TTL:      SessionTTL,
		Behavior: api.SessionBehaviorDelete,
	}, nil)
	if err != nil {
		return errors.Wrap(err, "failed to create consul session")
	}

	acquired, _, err := r.client.KV().Acquire(&api.KVPair{
		Key:     makeKey(item.SID),
		Value:   item.ToJSON(),
		Session: sessionID,
	}, nil)
	if err != nil {
		return errors.Wrap(err, "failed to create consul key for db session")
	}

	if !acquired {
		return fmt.Errorf("failed to acquire lock on session")
	}

	return err
}

// Get gets a session. Returns ErrNotFound if not found
func (r *ConsulRegistry) Get(sid string) (*RegItem, error) {

	kv, _, err := r.client.KV().Get(makeKey(sid), nil)
	if err != nil {
		return nil, errors.Wrap(err, "failed to get session")
	}
	if kv == nil {
		return nil, ErrNotFound
	}
	var item RegItem
	if err = util.FromJSON(kv.Value, &item); err != nil {
		return nil, errors.Wrap(err, "failed to parse registry item data")
	}
	return &item, nil
}

// Del deletes a session. Returns nil if the session was
// deleted successfully or it doesn't exists
func (r *ConsulRegistry) Del(sid string) (err error) {

	if !govalidator.IsUUIDv4(sid) {
		return fmt.Errorf("session id must be a UUID")
	}

	_, err = r.client.Session().Destroy(sid, nil)
	if err != nil {
		return errors.Wrap(err, "failed to delete registry item")
	}

	_, err = r.client.KV().Delete(makeKey(sid), nil)
	if err != nil {
		return errors.Wrap(err, "failed to delete session key")
	}

	return
}
<|endoftext|>


<commit_before>package adyen

/**********
* Payment *
**********/

// One-click functionality gives the shopper the option to store their payment details with the merchant, within the Adyen environment.
//
// In this type of transaction, the shopper needs to enter the CVC code for the transaction to get through.
//
// Link: https://docs.adyen.com/developers/api-reference/payments-api#recurring
const (
	RecurringPaymentOneClick              = "ONECLICK"
	RecurringPaymentRecurring             = "RECURRING"
	ShopperInteractionContAuth            = "ContAuth"
	SelectRecurringDetailReferenceLatests = "LATEST"
)

// AuthoriseEncrypted structure for Authorisation request (with encrypted card information)
//
// Link - https://docs.adyen.com/developers/api-reference/payments-api#paymentrequest
type AuthoriseEncrypted struct {
	AdditionalData                   *AdditionalData `json:"additionalData,omitempty"`
	Amount                           *Amount         `json:"amount"`
	BillingAddress                   *Address        `json:"billingAddress,omitempty"`
	DeliveryAddress                  *Address        `json:"deliveryAddress,omitempty"`
	Reference                        string          `json:"reference"`
	MerchantAccount                  string          `json:"merchantAccount"`
	ShopperReference                 string          `json:"shopperReference,omitempty"` // Mandatory for recurring payment
	Recurring                        *Recurring      `json:"recurring,omitempty"`
	ShopperEmail                     string          `json:"shopperEmail,omitempty"`
	ShopperInteraction               string          `json:"shopperInteraction,omitempty"`
	ShopperIP                        string          `json:"shopperIP,omitempty"`
	ShopperLocale                    string          `json:"shopperLocale,omitempty"`
	ShopperName                      *Name           `json:"shopperName,omitempty"`
	SelectedRecurringDetailReference string          `json:"selectedRecurringDetailReference,omitempty"`
	BrowserInfo                      *BrowserInfo    `json:"browserInfo,omitempty"` // Required for a 3DS process
	CaptureDelayHours                *int            `json:"captureDelayHours,omitempty"`
}

// Authorise structure for Authorisation request (card is not encrypted)
//
// Link - https://docs.adyen.com/developers/api-reference/payments-api#paymentrequest
type Authorise struct {
	AdditionalData                   *AdditionalData `json:"additionalData,omitempty"`
	Card                             *Card           `json:"card,omitempty"`
	Amount                           *Amount         `json:"amount"`
	BillingAddress                   *Address        `json:"billingAddress,omitempty"`
	DeliveryAddress                  *Address        `json:"deliveryAddress,omitempty"`
	Reference                        string          `json:"reference"`
	MerchantAccount                  string          `json:"merchantAccount"`
	ShopperReference                 string          `json:"shopperReference,omitempty"` // Mandatory for recurring payment
	Recurring                        *Recurring      `json:"recurring,omitempty"`
	ShopperEmail                     string          `json:"shopperEmail,omitempty"`
	ShopperInteraction               string          `json:"shopperInteraction,omitempty"`
	ShopperIP                        string          `json:"shopperIP,omitempty"`
	ShopperLocale                    string          `json:"shopperLocale,omitempty"`
	ShopperName                      *Name           `json:"shopperName,omitempty"`
	SelectedRecurringDetailReference string          `json:"selectedRecurringDetailReference,omitempty"`
	BrowserInfo                      *BrowserInfo    `json:"browserInfo,omitempty"` // Required for a 3DS process
	CaptureDelayHours                *int            `json:"captureDelayHours,omitempty"`
}

// AuthoriseResponse is a response structure for Adyen
//
// Link - https://docs.adyen.com/developers/api-reference/payments-api#paymentresult
type AuthoriseResponse struct {
	PspReference   string          `json:"pspReference"`
	ResultCode     string          `json:"resultCode"`
	AuthCode       string          `json:"authCode"`
	RefusalReason  string          `json:"refusalReason"`
	IssuerURL      string          `json:"issuerUrl"`
	MD             string          `json:"md"`
	PaRequest      string          `json:"paRepquest"`
	FraudResult    *FraudResult    `json:"fraudResult,omitempty"`
	AdditionalData *AdditionalData `json:"additionalData,omitempty"`
}

// AdditionalData stores encrypted information about customer's credit card
type AdditionalData struct {
	Content                           string      `json:"card.encrypted.json,omitempty"`
	AliasType                         string      `json:"aliasType,omitempty"`
	Alias                             string      `json:"alias,omitempty"`
	ExpiryDate                        string      `json:"expiryDate,omitempty"`
	CardBin                           string      `json:"cardBin,omitempty"`
	CardSummary                       string      `json:"cardSummary,omitempty"`
	CardHolderName                    string      `json:"cardHolderName,omitempty"`
	PaymentMethod                     string      `json:"paymentMethod,omitempty"`
	CardPaymentMethod                 string      `json:"cardPaymentMethod,omitempty"`
	CardIssuingCountry                string      `json:"cardIssuingCountry,omitempty"`
	RecurringDetailReference          string      `json:"recurring.recurringDetailReference,omitempty"`
	ExecuteThreeD                     *StringBool `json:"executeThreeD,omitempty"`
	FundingSource                     string      `json:"fundingSource,omitempty"`
	CustomRoutingFlag                 string      `json:"customRoutingFlag,omitempty"`
	RequestedTestAcquirerResponseCode int         `json:"RequestedTestAcquirerResponseCode,omitempty"` //Used for trigger error from adyen
}

// BrowserInfo hold information on the user browser
type BrowserInfo struct {
	AcceptHeader string `json:"acceptHeader"`
	UserAgent    string `json:"userAgent"`
}

// Recurring hold the behavior for a future payment : could be ONECLICK or RECURRING
type Recurring struct {
	Contract string `json:"contract"`
}

// FraudResult hold the fraud score of transaction
type FraudResult struct {
	AccountScore int64    `json:"accountScore,omitempty"`
	Results      []Result `json:"results,omitempty"`
}

// Result hold the fraud score detail
type Result struct {
	FraudCheckResult *FraudCheckResult `json:"FraudCheckResult,omitempty"`
}

// FraudCheckResult hold information of fraud score detail
type FraudCheckResult struct {
	AccountScore int    `json:"accountScore,omitempty"`
	CheckID      int    `json:"checkId,omitempty"`
	Name         string `json:"name,omitempty"`
}

/*************
* Payment 3D *
*************/

// Authorise3D structure for Authorisation request (card is not encrypted)
//
// https://docs.adyen.com/developers/api-reference/payments-api#paymentrequest3d
type Authorise3D struct {
	BillingAddress  *Address     `json:"billingAddress,omitempty"`
	DeliveryAddress *Address     `json:"deliveryAddress,omitempty"`
	MD              string       `json:"md"`
	MerchantAccount string       `json:"merchantAccount"`
	BrowserInfo     *BrowserInfo `json:"browserInfo"`
	PaResponse      string       `json:"paResponse"`
	ShopperEmail    string       `json:"shopperEmail,omitempty"`
	ShopperIP       string       `json:"shopperIP,omitempty"`
	ShopperLocale   string       `json:"shopperLocale,omitempty"`
	ShopperName     *Name        `json:"shopperName,omitempty"`
}

/*******************
* Directory lookup *
*******************/

// DirectoryLookupRequest - get list of available payment methods based on skin, country and order details
//
// Description - https://docs.adyen.com/developers/api-reference/hosted-payment-pages-api#directoryrequest
// CountryCode could be used to test local payment methods, if client's IP is from different country
type DirectoryLookupRequest struct {
	CurrencyCode      string `url:"currencyCode"`
	MerchantAccount   string `url:"merchantAccount"`
	PaymentAmount     int    `url:"paymentAmount"`
	SkinCode          string `url:"skinCode"`
	MerchantReference string `url:"merchantReference"`
	SessionsValidity  string `url:"sessionValidity"`
	MerchantSig       string `url:"merchantSig"`
	CountryCode       string `url:"countryCode"`
	ShipBeforeDate    string `url:"shipBeforeDate"`
}

// DirectoryLookupResponse - api response for DirectoryLookupRequest
//
// Description - https://docs.adyen.com/developers/api-reference/hosted-payment-pages-api#directoryresponse
type DirectoryLookupResponse struct {
	PaymentMethods []PaymentMethod `json:"paymentMethods"`
}

// PaymentMethod - structure for single payment method in directory look up response
//
// Part of DirectoryLookupResponse
type PaymentMethod struct {
	BrandCode string   `json:"brandCode"`
	Name      string   `json:"name"`
	Logos     logos    `json:"logos"`
	Issuers   []issuer `json:"issuers"`
}

// logos - payment method logos
//
// Part of DirectoryLookupResponse
type logos struct {
	Normal string `json:"normal"`
	Small  string `json:"small"`
	Tiny   string `json:"tiny"`
}

// issuer - bank issuer type
//
// Part of DirectoryLookupResponse
type issuer struct {
	IssuerID string `json:"issuerId"`
	Name     string `json:"name"`
}

/***********
* Skip HPP *
***********/

// SkipHppRequest contains data that would be used to create Adyen HPP redirect URL
//
// Link: https://docs.adyen.com/developers/ecommerce-integration/local-payment-methods
//
// Request description: https://docs.adyen.com/developers/api-reference/hosted-payment-pages-api#skipdetailsrequest
type SkipHppRequest struct {
	MerchantReference string `url:"merchantReference"`
	PaymentAmount     int    `url:"paymentAmount"`
	CurrencyCode      string `url:"currencyCode"`
	ShipBeforeDate    string `url:"shipBeforeDate"`
	SkinCode          string `url:"skinCode"`
	MerchantAccount   string `url:"merchantAccount"`
	ShopperLocale     string `url:"shopperLocale"`
	SessionsValidity  string `url:"sessionValidity"`
	MerchantSig       string `url:"merchantSig"`
	CountryCode       string `url:"countryCode"`
	BrandCode         string `url:"brandCode"`
	IssuerID          string `url:"issuerId"`
}
<commit_msg>Typos<commit_after>package adyen

/**********
* Payment *
**********/

// One-click functionality gives the shopper the option to store their payment details with the merchant, within the Adyen environment.
//
// In this type of transaction, the shopper needs to enter the CVC code for the transaction to get through.
//
// Link: https://docs.adyen.com/developers/api-reference/payments-api#recurring
const (
	RecurringPaymentOneClick              = "ONECLICK"
	RecurringPaymentRecurring             = "RECURRING"
	ShopperInteractionContAuth            = "ContAuth"
	SelectRecurringDetailReferenceLatests = "LATEST"
)

// AuthoriseEncrypted structure for Authorisation request (with encrypted card information)
//
// Link - https://docs.adyen.com/developers/api-reference/payments-api#paymentrequest
type AuthoriseEncrypted struct {
	AdditionalData                   *AdditionalData `json:"additionalData,omitempty"`
	Amount                           *Amount         `json:"amount"`
	BillingAddress                   *Address        `json:"billingAddress,omitempty"`
	DeliveryAddress                  *Address        `json:"deliveryAddress,omitempty"`
	Reference                        string          `json:"reference"`
	MerchantAccount                  string          `json:"merchantAccount"`
	ShopperReference                 string          `json:"shopperReference,omitempty"` // Mandatory for recurring payment
	Recurring                        *Recurring      `json:"recurring,omitempty"`
	ShopperEmail                     string          `json:"shopperEmail,omitempty"`
	ShopperInteraction               string          `json:"shopperInteraction,omitempty"`
	ShopperIP                        string          `json:"shopperIP,omitempty"`
	ShopperLocale                    string          `json:"shopperLocale,omitempty"`
	ShopperName                      *Name           `json:"shopperName,omitempty"`
	SelectedRecurringDetailReference string          `json:"selectedRecurringDetailReference,omitempty"`
	BrowserInfo                      *BrowserInfo    `json:"browserInfo,omitempty"` // Required for a 3DS process
	CaptureDelayHours                *int            `json:"captureDelayHours,omitempty"`
}

// Authorise structure for Authorisation request (card is not encrypted)
//
// Link - https://docs.adyen.com/developers/api-reference/payments-api#paymentrequest
type Authorise struct {
	AdditionalData                   *AdditionalData `json:"additionalData,omitempty"`
	Card                             *Card           `json:"card,omitempty"`
	Amount                           *Amount         `json:"amount"`
	BillingAddress                   *Address        `json:"billingAddress,omitempty"`
	DeliveryAddress                  *Address        `json:"deliveryAddress,omitempty"`
	Reference                        string          `json:"reference"`
	MerchantAccount                  string          `json:"merchantAccount"`
	ShopperReference                 string          `json:"shopperReference,omitempty"` // Mandatory for recurring payment
	Recurring                        *Recurring      `json:"recurring,omitempty"`
	ShopperEmail                     string          `json:"shopperEmail,omitempty"`
	ShopperInteraction               string          `json:"shopperInteraction,omitempty"`
	ShopperIP                        string          `json:"shopperIP,omitempty"`
	ShopperLocale                    string          `json:"shopperLocale,omitempty"`
	ShopperName                      *Name           `json:"shopperName,omitempty"`
	SelectedRecurringDetailReference string          `json:"selectedRecurringDetailReference,omitempty"`
	BrowserInfo                      *BrowserInfo    `json:"browserInfo,omitempty"` // Required for a 3DS process
	CaptureDelayHours                *int            `json:"captureDelayHours,omitempty"`
}

// AuthoriseResponse is a response structure for Adyen
//
// Link - https://docs.adyen.com/developers/api-reference/payments-api#paymentresult
type AuthoriseResponse struct {
	PspReference   string          `json:"pspReference"`
	ResultCode     string          `json:"resultCode"`
	AuthCode       string          `json:"authCode"`
	RefusalReason  string          `json:"refusalReason"`
	IssuerURL      string          `json:"issuerUrl"`
	MD             string          `json:"md"`
	PaRequest      string          `json:"paRequest"`
	FraudResult    *FraudResult    `json:"fraudResult,omitempty"`
	AdditionalData *AdditionalData `json:"additionalData,omitempty"`
}

// AdditionalData stores encrypted information about customer's credit card
type AdditionalData struct {
	Content                           string      `json:"card.encrypted.json,omitempty"`
	AliasType                         string      `json:"aliasType,omitempty"`
	Alias                             string      `json:"alias,omitempty"`
	ExpiryDate                        string      `json:"expiryDate,omitempty"`
	CardBin                           string      `json:"cardBin,omitempty"`
	CardSummary                       string      `json:"cardSummary,omitempty"`
	CardHolderName                    string      `json:"cardHolderName,omitempty"`
	PaymentMethod                     string      `json:"paymentMethod,omitempty"`
	CardPaymentMethod                 string      `json:"cardPaymentMethod,omitempty"`
	CardIssuingCountry                string      `json:"cardIssuingCountry,omitempty"`
	RecurringDetailReference          string      `json:"recurring.recurringDetailReference,omitempty"`
	ExecuteThreeD                     *StringBool `json:"executeThreeD,omitempty"`
	FundingSource                     string      `json:"fundingSource,omitempty"`
	CustomRoutingFlag                 string      `json:"customRoutingFlag,omitempty"`
	RequestedTestAcquirerResponseCode int         `json:"RequestedTestAcquirerResponseCode,omitempty"` //Used for trigger error from adyen
}

// BrowserInfo hold information on the user browser
type BrowserInfo struct {
	AcceptHeader string `json:"acceptHeader"`
	UserAgent    string `json:"userAgent"`
}

// Recurring hold the behavior for a future payment : could be ONECLICK or RECURRING
type Recurring struct {
	Contract string `json:"contract"`
}

// FraudResult hold the fraud score of transaction
type FraudResult struct {
	AccountScore int64    `json:"accountScore,omitempty"`
	Results      []Result `json:"results,omitempty"`
}

// Result hold the fraud score detail
type Result struct {
	FraudCheckResult *FraudCheckResult `json:"FraudCheckResult,omitempty"`
}

// FraudCheckResult hold information of fraud score detail
type FraudCheckResult struct {
	AccountScore int    `json:"accountScore,omitempty"`
	CheckID      int    `json:"checkId,omitempty"`
	Name         string `json:"name,omitempty"`
}

/*************
* Payment 3D *
*************/

// Authorise3D structure for Authorisation request (card is not encrypted)
//
// https://docs.adyen.com/developers/api-reference/payments-api#paymentrequest3d
type Authorise3D struct {
	BillingAddress  *Address     `json:"billingAddress,omitempty"`
	DeliveryAddress *Address     `json:"deliveryAddress,omitempty"`
	MD              string       `json:"md"`
	MerchantAccount string       `json:"merchantAccount"`
	BrowserInfo     *BrowserInfo `json:"browserInfo"`
	PaResponse      string       `json:"paResponse"`
	ShopperEmail    string       `json:"shopperEmail,omitempty"`
	ShopperIP       string       `json:"shopperIP,omitempty"`
	ShopperLocale   string       `json:"shopperLocale,omitempty"`
	ShopperName     *Name        `json:"shopperName,omitempty"`
}

/*******************
* Directory lookup *
*******************/

// DirectoryLookupRequest - get list of available payment methods based on skin, country and order details
//
// Description - https://docs.adyen.com/developers/api-reference/hosted-payment-pages-api#directoryrequest
// CountryCode could be used to test local payment methods, if client's IP is from different country
type DirectoryLookupRequest struct {
	CurrencyCode      string `url:"currencyCode"`
	MerchantAccount   string `url:"merchantAccount"`
	PaymentAmount     int    `url:"paymentAmount"`
	SkinCode          string `url:"skinCode"`
	MerchantReference string `url:"merchantReference"`
	SessionsValidity  string `url:"sessionValidity"`
	MerchantSig       string `url:"merchantSig"`
	CountryCode       string `url:"countryCode"`
	ShipBeforeDate    string `url:"shipBeforeDate"`
}

// DirectoryLookupResponse - api response for DirectoryLookupRequest
//
// Description - https://docs.adyen.com/developers/api-reference/hosted-payment-pages-api#directoryresponse
type DirectoryLookupResponse struct {
	PaymentMethods []PaymentMethod `json:"paymentMethods"`
}

// PaymentMethod - structure for single payment method in directory look up response
//
// Part of DirectoryLookupResponse
type PaymentMethod struct {
	BrandCode string   `json:"brandCode"`
	Name      string   `json:"name"`
	Logos     logos    `json:"logos"`
	Issuers   []issuer `json:"issuers"`
}

// logos - payment method logos
//
// Part of DirectoryLookupResponse
type logos struct {
	Normal string `json:"normal"`
	Small  string `json:"small"`
	Tiny   string `json:"tiny"`
}

// issuer - bank issuer type
//
// Part of DirectoryLookupResponse
type issuer struct {
	IssuerID string `json:"issuerId"`
	Name     string `json:"name"`
}

/***********
* Skip HPP *
***********/

// SkipHppRequest contains data that would be used to create Adyen HPP redirect URL
//
// Link: https://docs.adyen.com/developers/ecommerce-integration/local-payment-methods
//
// Request description: https://docs.adyen.com/developers/api-reference/hosted-payment-pages-api#skipdetailsrequest
type SkipHppRequest struct {
	MerchantReference string `url:"merchantReference"`
	PaymentAmount     int    `url:"paymentAmount"`
	CurrencyCode      string `url:"currencyCode"`
	ShipBeforeDate    string `url:"shipBeforeDate"`
	SkinCode          string `url:"skinCode"`
	MerchantAccount   string `url:"merchantAccount"`
	ShopperLocale     string `url:"shopperLocale"`
	SessionsValidity  string `url:"sessionValidity"`
	MerchantSig       string `url:"merchantSig"`
	CountryCode       string `url:"countryCode"`
	BrandCode         string `url:"brandCode"`
	IssuerID          string `url:"issuerId"`
}
<|endoftext|>


<commit_before>package wallet

import (
	"bytes"
	"crypto/rand"
	"errors"
	"time"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"
	"github.com/NebulousLabs/bolt"
)

var (
	errAlreadyUnlocked   = errors.New("wallet has already been unlocked")
	errReencrypt         = errors.New("wallet is already encrypted, cannot encrypt again")
	errUnencryptedWallet = errors.New("wallet has not been encrypted yet")

	// verificationPlaintext is the plaintext used to verify encryption keys.
	// By storing the corresponding ciphertext for a given key, we can later
	// verify that a key is correct by using it to decrypt the ciphertext and
	// comparing the result to verificationPlaintext.
	verificationPlaintext = make([]byte, 32)
)

// uidEncryptionKey creates an encryption key that is used to decrypt a
// specific key file.
func uidEncryptionKey(masterKey crypto.TwofishKey, uid uniqueID) crypto.TwofishKey {
	return crypto.TwofishKey(crypto.HashAll(masterKey, uid))
}

// verifyEncryption verifies that key properly decrypts the ciphertext to a
// preset plaintext.
func verifyEncryption(key crypto.TwofishKey, encrypted crypto.Ciphertext) error {
	verification, err := key.DecryptBytes(encrypted)
	if err != nil {
		return modules.ErrBadEncryptionKey
	}
	if !bytes.Equal(verificationPlaintext, verification) {
		return modules.ErrBadEncryptionKey
	}
	return nil
}

// checkMasterKey verifies that the masterKey is the key used to encrypt the wallet.
func checkMasterKey(tx *bolt.Tx, masterKey crypto.TwofishKey) error {
	uk := uidEncryptionKey(masterKey, dbGetWalletUID(tx))
	encryptedVerification := tx.Bucket(bucketWallet).Get(keyEncryptionVerification)
	return verifyEncryption(uk, encryptedVerification)
}

// initEncryption initializes and encrypts the primary SeedFile.
func (w *Wallet) initEncryption(masterKey crypto.TwofishKey, seed modules.Seed) (modules.Seed, error) {
	err := w.db.Update(func(tx *bolt.Tx) error {
		wb := tx.Bucket(bucketWallet)
		// Check if the wallet encryption key has already been set.
		if wb.Get(keyEncryptionVerification) != nil {
			return errReencrypt
		}

		// create a seedFile for the seed
		sf, err := createSeedFile(masterKey, seed)
		if err != nil {
			return err
		}

		// set this as the primary seedFile
		err = wb.Put(keyPrimarySeedFile, encoding.Marshal(sf))
		if err != nil {
			return err
		}
		err = wb.Put(keyPrimarySeedProgress, encoding.Marshal(uint64(0)))
		if err != nil {
			return err
		}

		// Establish the encryption verification using the masterKey. After this
		// point, the wallet is encrypted.
		uk := uidEncryptionKey(masterKey, dbGetWalletUID(tx))
		verification, err := uk.EncryptBytes(verificationPlaintext)
		if err != nil {
			return err
		}
		err = wb.Put(keyEncryptionVerification, verification[:])
		if err != nil {
			return err
		}
		return nil
	})
	if err != nil {
		return modules.Seed{}, err
	}

	// on future startups, this field will be set by w.initPersist
	w.encrypted = true

	return seed, nil
}

// managedUnlock loads all of the encrypted file structures into wallet memory. Even
// after loading, the structures are kept encrypted, but some data such as
// addresses are decrypted so that the wallet knows what to track.
func (w *Wallet) managedUnlock(masterKey crypto.TwofishKey) error {
	w.mu.RLock()
	unlocked := w.unlocked
	encrypted := w.encrypted
	w.mu.RUnlock()
	if unlocked {
		return errAlreadyUnlocked
	} else if !encrypted {
		return errUnencryptedWallet
	}

	// Load db objects into memory.
	var lastChange modules.ConsensusChangeID
	var primarySeedFile seedFile
	var primarySeedProgress uint64
	var auxiliarySeedFiles []seedFile
	var unseededKeyFiles []spendableKeyFile
	err := w.db.View(func(tx *bolt.Tx) error {
		// verify masterKey
		err := checkMasterKey(tx, masterKey)
		if err != nil {
			return err
		}

		// lastChange
		lastChange = dbGetConsensusChangeID(tx)

		// primarySeedFile + primarySeedProgress
		err = encoding.Unmarshal(tx.Bucket(bucketWallet).Get(keyPrimarySeedFile), &primarySeedFile)
		if err != nil {
			return err
		}
		err = encoding.Unmarshal(tx.Bucket(bucketWallet).Get(keyPrimarySeedProgress), &primarySeedProgress)
		if err != nil {
			return err
		}

		// auxiliarySeedFiles
		err = tx.Bucket(bucketSeedFiles).ForEach(func(_, sfBytes []byte) error {
			var sf seedFile
			err := encoding.Unmarshal(sfBytes, &sf)
			if err != nil {
				return err
			}
			auxiliarySeedFiles = append(auxiliarySeedFiles, sf)
			return nil
		})
		if err != nil {
			return err
		}

		// unseededKeyFiles
		err = tx.Bucket(bucketSpendableKeyFiles).ForEach(func(_, ukfBytes []byte) error {
			var ukf spendableKeyFile
			err := encoding.Unmarshal(ukfBytes, &ukf)
			if err != nil {
				return err
			}
			unseededKeyFiles = append(unseededKeyFiles, ukf)
			return nil
		})
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		return err
	}

	// Decrypt + load keys.
	err = func() error {
		w.mu.Lock()
		defer w.mu.Unlock()

		// primarySeedFile
		primarySeed, err := decryptSeedFile(masterKey, primarySeedFile)
		if err != nil {
			return err
		}
		w.integrateSeed(primarySeed, primarySeedProgress)
		w.primarySeed = primarySeed

		// auxiliarySeedFiles
		for _, sf := range auxiliarySeedFiles {
			auxSeed, err := decryptSeedFile(masterKey, sf)
			if err != nil {
				return err
			}
			w.integrateSeed(auxSeed, modules.PublicKeysPerSeed)
			w.seeds = append(w.seeds, auxSeed)
		}

		// unseededKeyFiles
		for _, uk := range unseededKeyFiles {
			sk, err := decryptSpendableKeyFile(masterKey, uk)
			if err != nil {
				return err
			}
			w.integrateSpendableKey(masterKey, sk)
		}
		return nil
	}()
	if err != nil {
		return err
	}

	// Subscribe to the consensus set if this is the first unlock for the
	// wallet object.
	w.mu.RLock()
	subscribed := w.subscribed
	w.mu.RUnlock()
	if !subscribed {
		err = w.cs.ConsensusSetSubscribe(w, lastChange)
		if err == modules.ErrInvalidConsensusChangeID {
			// something went wrong; resubscribe from the beginning and spawn a
			// goroutine to display rescan progress
			if build.Release != "testing" {
				go func() {
					println("Rescanning consensus set...")
					for range time.Tick(time.Second * 3) {
						w.mu.RLock()
						var height types.BlockHeight
						w.db.View(func(tx *bolt.Tx) error {
							var err error
							height, err = dbGetConsensusHeight(tx)
							return err
						})
						done := w.subscribed
						w.mu.RUnlock()
						if done {
							println("\nDone!")
							break
						}
						print("\rScanned to height ", height, "...")
					}
				}()
			}
			// set the db entries for consensus change and consensus height
			err = w.db.Update(func(tx *bolt.Tx) error {
				err := dbPutConsensusChangeID(tx, modules.ConsensusChangeBeginning)
				if err != nil {
					return err
				}
				return dbPutConsensusHeight(tx, 0)
			})
			if err != nil {
				return errors.New("failed to reset db during rescan: " + err.Error())
			}
			err = w.cs.ConsensusSetSubscribe(w, modules.ConsensusChangeBeginning)
		}
		if err != nil {
			return errors.New("wallet subscription failed: " + err.Error())
		}
		w.tpool.TransactionPoolSubscribe(w)
		w.mu.Lock()
		w.subscribed = true
		w.mu.Unlock()
	}

	w.mu.Lock()
	w.unlocked = true
	w.mu.Unlock()
	return nil
}

// wipeSecrets erases all of the seeds and secret keys in the wallet.
func (w *Wallet) wipeSecrets() {
	// 'for i := range' must be used to prevent copies of secret data from
	// being made.
	for i := range w.keys {
		for j := range w.keys[i].SecretKeys {
			crypto.SecureWipe(w.keys[i].SecretKeys[j][:])
		}
	}
	for i := range w.seeds {
		crypto.SecureWipe(w.seeds[i][:])
	}
	crypto.SecureWipe(w.primarySeed[:])
	w.seeds = w.seeds[:0]
}

// Encrypted returns whether or not the wallet has been encrypted.
func (w *Wallet) Encrypted() bool {
	w.mu.Lock()
	defer w.mu.Unlock()
	if build.DEBUG && w.unlocked && !w.encrypted {
		panic("wallet is both unlocked and unencrypted")
	}
	return w.encrypted
}

// Encrypt will create a primary seed for the wallet and encrypt it using
// masterKey. If masterKey is blank, then the hash of the primary seed will be
// used instead. The wallet will still be locked after Encrypt is called.
//
// Encrypt can only be called once throughout the life of the wallet, and will
// return an error on subsequent calls (even after restarting the wallet). To
// reset the wallet, the wallet files must be moved to a different directory
// or deleted.
func (w *Wallet) Encrypt(masterKey crypto.TwofishKey) (modules.Seed, error) {
	if err := w.tg.Add(); err != nil {
		return modules.Seed{}, err
	}
	defer w.tg.Done()
	w.mu.Lock()
	defer w.mu.Unlock()

	// Create a random seed.
	var seed modules.Seed
	_, err := rand.Read(seed[:])
	if err != nil {
		return modules.Seed{}, err
	}

	// If masterKey is blank, use the hash of the seed.
	if masterKey == (crypto.TwofishKey{}) {
		masterKey = crypto.TwofishKey(crypto.HashObject(seed))
	}

	return w.initEncryption(masterKey, seed)
}

// InitFromSeed functions like Init, but using a specified seed. Unlike Init,
// the blockchain will be scanned to determine the seed's progress. For this
// reason, InitFromSeed should not be called until the blockchain is fully
// synced.
func (w *Wallet) InitFromSeed(masterKey crypto.TwofishKey, seed modules.Seed) error {
	if err := w.tg.Add(); err != nil {
		return err
	}
	defer w.tg.Done()

	if !w.cs.Synced() {
		return errors.New("cannot init from seed until blockchain is synced")
	}

	// If masterKey is blank, use the hash of the seed.
	if masterKey == (crypto.TwofishKey{}) {
		masterKey = crypto.TwofishKey(crypto.HashObject(seed))
	}

	w.mu.Lock()
	defer w.mu.Unlock()
	if _, err := w.initEncryption(masterKey, seed); err != nil {
		return err
	}

	// estimate the primarySeedProgress by scanning the blockchain
	s := newSeedScanner(seed)
	if err := s.scan(w.cs); err != nil {
		return err
	}
	progress := s.largestIndexSeen + (s.largestIndexSeen / 10) // add 10% to be safe
	// set primarySeedProgress
	return w.db.Update(func(tx *bolt.Tx) error {
		return dbPutPrimarySeedProgress(tx, uint64(progress))
	})
}

// Unlocked indicates whether the wallet is locked or unlocked.
func (w *Wallet) Unlocked() bool {
	w.mu.RLock()
	defer w.mu.RUnlock()
	return w.unlocked
}

// Lock will erase all keys from memory and prevent the wallet from spending
// coins until it is unlocked.
func (w *Wallet) Lock() error {
	w.mu.Lock()
	defer w.mu.Unlock()
	if !w.unlocked {
		return modules.ErrLockedWallet
	}
	w.log.Println("INFO: Locking wallet.")

	// Wipe all of the seeds and secret keys. They will be replaced upon
	// calling 'Unlock' again. Note that since the public keys are not wiped,
	// we can continue processing blocks.
	w.wipeSecrets()
	w.unlocked = false
	return nil
}

// Unlock will decrypt the wallet seed and load all of the addresses into
// memory.
func (w *Wallet) Unlock(masterKey crypto.TwofishKey) error {
	// By having the wallet's ThreadGroup track the Unlock method, we ensure
	// that Unlock will never unlock the wallet once the ThreadGroup has been
	// stopped. Without this precaution, the wallet's Close method would be
	// unsafe because it would theoretically be possible for another function
	// to Unlock the wallet in the short interval after Close calls w.Lock
	// and before Close calls w.mu.Lock.
	if err := w.tg.Add(); err != nil {
		return err
	}
	defer w.tg.Done()
	w.log.Println("INFO: Unlocking wallet.")

	// Initialize all of the keys in the wallet under a lock. While holding the
	// lock, also grab the subscriber status.
	return w.managedUnlock(masterKey)
}
<commit_msg>only hold lock during initEncryption<commit_after>package wallet

import (
	"bytes"
	"crypto/rand"
	"errors"
	"time"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"
	"github.com/NebulousLabs/bolt"
)

var (
	errAlreadyUnlocked   = errors.New("wallet has already been unlocked")
	errReencrypt         = errors.New("wallet is already encrypted, cannot encrypt again")
	errUnencryptedWallet = errors.New("wallet has not been encrypted yet")

	// verificationPlaintext is the plaintext used to verify encryption keys.
	// By storing the corresponding ciphertext for a given key, we can later
	// verify that a key is correct by using it to decrypt the ciphertext and
	// comparing the result to verificationPlaintext.
	verificationPlaintext = make([]byte, 32)
)

// uidEncryptionKey creates an encryption key that is used to decrypt a
// specific key file.
func uidEncryptionKey(masterKey crypto.TwofishKey, uid uniqueID) crypto.TwofishKey {
	return crypto.TwofishKey(crypto.HashAll(masterKey, uid))
}

// verifyEncryption verifies that key properly decrypts the ciphertext to a
// preset plaintext.
func verifyEncryption(key crypto.TwofishKey, encrypted crypto.Ciphertext) error {
	verification, err := key.DecryptBytes(encrypted)
	if err != nil {
		return modules.ErrBadEncryptionKey
	}
	if !bytes.Equal(verificationPlaintext, verification) {
		return modules.ErrBadEncryptionKey
	}
	return nil
}

// checkMasterKey verifies that the masterKey is the key used to encrypt the wallet.
func checkMasterKey(tx *bolt.Tx, masterKey crypto.TwofishKey) error {
	uk := uidEncryptionKey(masterKey, dbGetWalletUID(tx))
	encryptedVerification := tx.Bucket(bucketWallet).Get(keyEncryptionVerification)
	return verifyEncryption(uk, encryptedVerification)
}

// initEncryption initializes and encrypts the primary SeedFile.
func (w *Wallet) initEncryption(masterKey crypto.TwofishKey, seed modules.Seed) (modules.Seed, error) {
	err := w.db.Update(func(tx *bolt.Tx) error {
		wb := tx.Bucket(bucketWallet)
		// Check if the wallet encryption key has already been set.
		if wb.Get(keyEncryptionVerification) != nil {
			return errReencrypt
		}

		// create a seedFile for the seed
		sf, err := createSeedFile(masterKey, seed)
		if err != nil {
			return err
		}

		// set this as the primary seedFile
		err = wb.Put(keyPrimarySeedFile, encoding.Marshal(sf))
		if err != nil {
			return err
		}
		err = wb.Put(keyPrimarySeedProgress, encoding.Marshal(uint64(0)))
		if err != nil {
			return err
		}

		// Establish the encryption verification using the masterKey. After this
		// point, the wallet is encrypted.
		uk := uidEncryptionKey(masterKey, dbGetWalletUID(tx))
		verification, err := uk.EncryptBytes(verificationPlaintext)
		if err != nil {
			return err
		}
		err = wb.Put(keyEncryptionVerification, verification[:])
		if err != nil {
			return err
		}
		return nil
	})
	if err != nil {
		return modules.Seed{}, err
	}

	// on future startups, this field will be set by w.initPersist
	w.encrypted = true

	return seed, nil
}

// managedUnlock loads all of the encrypted file structures into wallet memory. Even
// after loading, the structures are kept encrypted, but some data such as
// addresses are decrypted so that the wallet knows what to track.
func (w *Wallet) managedUnlock(masterKey crypto.TwofishKey) error {
	w.mu.RLock()
	unlocked := w.unlocked
	encrypted := w.encrypted
	w.mu.RUnlock()
	if unlocked {
		return errAlreadyUnlocked
	} else if !encrypted {
		return errUnencryptedWallet
	}

	// Load db objects into memory.
	var lastChange modules.ConsensusChangeID
	var primarySeedFile seedFile
	var primarySeedProgress uint64
	var auxiliarySeedFiles []seedFile
	var unseededKeyFiles []spendableKeyFile
	err := w.db.View(func(tx *bolt.Tx) error {
		// verify masterKey
		err := checkMasterKey(tx, masterKey)
		if err != nil {
			return err
		}

		// lastChange
		lastChange = dbGetConsensusChangeID(tx)

		// primarySeedFile + primarySeedProgress
		err = encoding.Unmarshal(tx.Bucket(bucketWallet).Get(keyPrimarySeedFile), &primarySeedFile)
		if err != nil {
			return err
		}
		err = encoding.Unmarshal(tx.Bucket(bucketWallet).Get(keyPrimarySeedProgress), &primarySeedProgress)
		if err != nil {
			return err
		}

		// auxiliarySeedFiles
		err = tx.Bucket(bucketSeedFiles).ForEach(func(_, sfBytes []byte) error {
			var sf seedFile
			err := encoding.Unmarshal(sfBytes, &sf)
			if err != nil {
				return err
			}
			auxiliarySeedFiles = append(auxiliarySeedFiles, sf)
			return nil
		})
		if err != nil {
			return err
		}

		// unseededKeyFiles
		err = tx.Bucket(bucketSpendableKeyFiles).ForEach(func(_, ukfBytes []byte) error {
			var ukf spendableKeyFile
			err := encoding.Unmarshal(ukfBytes, &ukf)
			if err != nil {
				return err
			}
			unseededKeyFiles = append(unseededKeyFiles, ukf)
			return nil
		})
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		return err
	}

	// Decrypt + load keys.
	err = func() error {
		w.mu.Lock()
		defer w.mu.Unlock()

		// primarySeedFile
		primarySeed, err := decryptSeedFile(masterKey, primarySeedFile)
		if err != nil {
			return err
		}
		w.integrateSeed(primarySeed, primarySeedProgress)
		w.primarySeed = primarySeed

		// auxiliarySeedFiles
		for _, sf := range auxiliarySeedFiles {
			auxSeed, err := decryptSeedFile(masterKey, sf)
			if err != nil {
				return err
			}
			w.integrateSeed(auxSeed, modules.PublicKeysPerSeed)
			w.seeds = append(w.seeds, auxSeed)
		}

		// unseededKeyFiles
		for _, uk := range unseededKeyFiles {
			sk, err := decryptSpendableKeyFile(masterKey, uk)
			if err != nil {
				return err
			}
			w.integrateSpendableKey(masterKey, sk)
		}
		return nil
	}()
	if err != nil {
		return err
	}

	// Subscribe to the consensus set if this is the first unlock for the
	// wallet object.
	w.mu.RLock()
	subscribed := w.subscribed
	w.mu.RUnlock()
	if !subscribed {
		err = w.cs.ConsensusSetSubscribe(w, lastChange)
		if err == modules.ErrInvalidConsensusChangeID {
			// something went wrong; resubscribe from the beginning and spawn a
			// goroutine to display rescan progress
			if build.Release != "testing" {
				go func() {
					println("Rescanning consensus set...")
					for range time.Tick(time.Second * 3) {
						w.mu.RLock()
						var height types.BlockHeight
						w.db.View(func(tx *bolt.Tx) error {
							var err error
							height, err = dbGetConsensusHeight(tx)
							return err
						})
						done := w.subscribed
						w.mu.RUnlock()
						if done {
							println("\nDone!")
							break
						}
						print("\rScanned to height ", height, "...")
					}
				}()
			}
			// set the db entries for consensus change and consensus height
			err = w.db.Update(func(tx *bolt.Tx) error {
				err := dbPutConsensusChangeID(tx, modules.ConsensusChangeBeginning)
				if err != nil {
					return err
				}
				return dbPutConsensusHeight(tx, 0)
			})
			if err != nil {
				return errors.New("failed to reset db during rescan: " + err.Error())
			}
			err = w.cs.ConsensusSetSubscribe(w, modules.ConsensusChangeBeginning)
		}
		if err != nil {
			return errors.New("wallet subscription failed: " + err.Error())
		}
		w.tpool.TransactionPoolSubscribe(w)
		w.mu.Lock()
		w.subscribed = true
		w.mu.Unlock()
	}

	w.mu.Lock()
	w.unlocked = true
	w.mu.Unlock()
	return nil
}

// wipeSecrets erases all of the seeds and secret keys in the wallet.
func (w *Wallet) wipeSecrets() {
	// 'for i := range' must be used to prevent copies of secret data from
	// being made.
	for i := range w.keys {
		for j := range w.keys[i].SecretKeys {
			crypto.SecureWipe(w.keys[i].SecretKeys[j][:])
		}
	}
	for i := range w.seeds {
		crypto.SecureWipe(w.seeds[i][:])
	}
	crypto.SecureWipe(w.primarySeed[:])
	w.seeds = w.seeds[:0]
}

// Encrypted returns whether or not the wallet has been encrypted.
func (w *Wallet) Encrypted() bool {
	w.mu.Lock()
	defer w.mu.Unlock()
	if build.DEBUG && w.unlocked && !w.encrypted {
		panic("wallet is both unlocked and unencrypted")
	}
	return w.encrypted
}

// Encrypt will create a primary seed for the wallet and encrypt it using
// masterKey. If masterKey is blank, then the hash of the primary seed will be
// used instead. The wallet will still be locked after Encrypt is called.
//
// Encrypt can only be called once throughout the life of the wallet, and will
// return an error on subsequent calls (even after restarting the wallet). To
// reset the wallet, the wallet files must be moved to a different directory
// or deleted.
func (w *Wallet) Encrypt(masterKey crypto.TwofishKey) (modules.Seed, error) {
	if err := w.tg.Add(); err != nil {
		return modules.Seed{}, err
	}
	defer w.tg.Done()
	w.mu.Lock()
	defer w.mu.Unlock()

	// Create a random seed.
	var seed modules.Seed
	_, err := rand.Read(seed[:])
	if err != nil {
		return modules.Seed{}, err
	}

	// If masterKey is blank, use the hash of the seed.
	if masterKey == (crypto.TwofishKey{}) {
		masterKey = crypto.TwofishKey(crypto.HashObject(seed))
	}

	return w.initEncryption(masterKey, seed)
}

// InitFromSeed functions like Init, but using a specified seed. Unlike Init,
// the blockchain will be scanned to determine the seed's progress. For this
// reason, InitFromSeed should not be called until the blockchain is fully
// synced.
func (w *Wallet) InitFromSeed(masterKey crypto.TwofishKey, seed modules.Seed) error {
	if err := w.tg.Add(); err != nil {
		return err
	}
	defer w.tg.Done()

	if !w.cs.Synced() {
		return errors.New("cannot init from seed until blockchain is synced")
	}

	// If masterKey is blank, use the hash of the seed.
	if masterKey == (crypto.TwofishKey{}) {
		masterKey = crypto.TwofishKey(crypto.HashObject(seed))
	}

	w.mu.Lock()
	if _, err := w.initEncryption(masterKey, seed); err != nil {
		w.mu.Unlock()
		return err
	}
	w.mu.Unlock()

	// estimate the primarySeedProgress by scanning the blockchain
	s := newSeedScanner(seed)
	if err := s.scan(w.cs); err != nil {
		return err
	}
	progress := s.largestIndexSeen + (s.largestIndexSeen / 10) // add 10% to be safe
	// set primarySeedProgress
	return w.db.Update(func(tx *bolt.Tx) error {
		return dbPutPrimarySeedProgress(tx, uint64(progress))
	})
}

// Unlocked indicates whether the wallet is locked or unlocked.
func (w *Wallet) Unlocked() bool {
	w.mu.RLock()
	defer w.mu.RUnlock()
	return w.unlocked
}

// Lock will erase all keys from memory and prevent the wallet from spending
// coins until it is unlocked.
func (w *Wallet) Lock() error {
	w.mu.Lock()
	defer w.mu.Unlock()
	if !w.unlocked {
		return modules.ErrLockedWallet
	}
	w.log.Println("INFO: Locking wallet.")

	// Wipe all of the seeds and secret keys. They will be replaced upon
	// calling 'Unlock' again. Note that since the public keys are not wiped,
	// we can continue processing blocks.
	w.wipeSecrets()
	w.unlocked = false
	return nil
}

// Unlock will decrypt the wallet seed and load all of the addresses into
// memory.
func (w *Wallet) Unlock(masterKey crypto.TwofishKey) error {
	// By having the wallet's ThreadGroup track the Unlock method, we ensure
	// that Unlock will never unlock the wallet once the ThreadGroup has been
	// stopped. Without this precaution, the wallet's Close method would be
	// unsafe because it would theoretically be possible for another function
	// to Unlock the wallet in the short interval after Close calls w.Lock
	// and before Close calls w.mu.Lock.
	if err := w.tg.Add(); err != nil {
		return err
	}
	defer w.tg.Done()
	w.log.Println("INFO: Unlocking wallet.")

	// Initialize all of the keys in the wallet under a lock. While holding the
	// lock, also grab the subscriber status.
	return w.managedUnlock(masterKey)
}
<|endoftext|>


<commit_before>package pdb

import (
	"bufio"
	"compress/gzip"
	"fmt"
	"io"
	"os"
	"path"
	"strconv"
	"strings"
)

// AminoThreeToOne is a map from three letter amino acids to their
// corresponding single letter representation.
var AminoThreeToOne = map[string]byte{
	"ALA": 'A', "ARG": 'R', "ASN": 'N', "ASP": 'D', "CYS": 'C',
	"GLU": 'E', "GLN": 'Q', "GLY": 'G', "HIS": 'H', "ILE": 'I',
	"LEU": 'L', "LYS": 'K', "MET": 'M', "PHE": 'F', "PRO": 'P',
	"SER": 'S', "THR": 'T', "TRP": 'W', "TYR": 'Y', "VAL": 'V',
	"SEC": 'U', "PYL": 'O',
	"UNK": 'X', "ACE": 'X', "NH2": 'X',
	"ASX": 'X', "GLX": 'X',
}

// AminoOneToThree is the reverse of AminoThreeToOne. It is created in
// this packages 'init' function.
var AminoOneToThree = map[byte]string{}

func init() {
	// Create a reverse map of AminoThreeToOne.
	for k, v := range AminoThreeToOne {
		AminoOneToThree[v] = k
	}
}

// Entry represents all information known about a particular PDB file (that
// has been implemented in this package).
//
// Currently, a PDB entry is simply a file path and a map of protein chains.
type Entry struct {
	Path           string
	IdCode         string
	Classification string
	Chains         []*Chain
}

// New creates a new PDB Entry from a file. If the file cannot be read, or there
// is an error parsing the PDB file, an error is returned.
//
// If the file name ends with ".gz", gzip decompression will be used.
func New(fileName string) (*Entry, error) {
	var reader io.Reader
	var err error

	reader, err = os.Open(fileName)
	if err != nil {
		return nil, err
	}

	// If the file is gzipped, use the gzip decompressor.
	if path.Ext(fileName) == ".gz" {
		reader, err = gzip.NewReader(reader)
		if err != nil {
			return nil, err
		}
	}

	entry := &Entry{
		Path:   fileName,
		Chains: make([]*Chain, 0),
	}

	// Now traverse each line, and process it according to the record name.
	// Note that it is imperative that we preserve the order of ATOM records
	// as we read them. We are currently trying to replicate Fragbag, and this
	// is what Fragbag does. (A more stable approach would require more
	// information from the PDB file; like differentiating models, since
	// sorting on ATOM serial number isn't good enough.)
	breader := bufio.NewReaderSize(reader, 1000)
	for {
		// We ignore 'isPrefix' here, since we never care about lines longer
		// than 1000 characters, which is the size of our buffer.
		line, _, err := breader.ReadLine()
		if err == io.EOF {
			break
		} else if err != nil {
			return nil, err
		}

		// The record name is always in the fix six columns.
		switch strings.TrimSpace(string(line[0:6])) {
		case "HEADER":
			if err := entry.parseHeader(line); err != nil {
				return nil, err
			}
		case "SEQRES":
			entry.parseSeqres(line)
		case "ATOM":
			entry.parseAtom(line)
		}
	}

	// If we didn't pick up any chains, this probably isn't a valid PDB file.
	if len(entry.Chains) == 0 {
		return nil, fmt.Errorf("The file '%s' does not appear to be a valid "+
			"PDB file.", fileName)
	}

	// If we couldn't find an Id code, inspect the base name of the file path.
	if len(entry.IdCode) == 0 {
		name := path.Base(fileName)
		switch {
		case len(name) >= 7 && name[0:3] == "pdb":
			entry.IdCode = name[3:7]
		case len(name) == 7: // cath
			entry.IdCode = name[0:4]
		}
	}

	return entry, nil
}

// Chain looks for the chain with identifier ident and returns it. 'nil' is
// returned if the chain could not be found.
func (e *Entry) Chain(ident byte) *Chain {
	for _, chain := range e.Chains {
		if chain.Ident == ident {
			return chain
		}
	}
	return nil
}

// OneChain returns a single chain in the PDB file. If there is more than one
// chain, OneChain will panic. This is convenient when you expect a PDB file to
// have only a single chain, but don't know the name.
func (e *Entry) OneChain() *Chain {
	if len(e.Chains) != 1 {
		panic(fmt.Sprintf("OneChain can only be called on PDB entries with "+
			"ONE chain. But the '%s' PDB entry has %d chains.",
			e.Path, len(e.Chains)))
	}
	return e.Chains[0]
}

// Name returns the base name of the path of this PDB entry.
func (e *Entry) Name() string {
	return path.Base(e.Path)
}

// String returns a list of all chains, their residue start/stop indices,
// and the amino acid sequence.
func (e *Entry) String() string {
	lines := make([]string, 0)
	for _, chain := range e.Chains {
		lines = append(lines, chain.String())
	}
	return strings.Join(lines, "\n")
}

// getOrMakeChain looks for a chain in the 'Chains' slice corresponding to the
// chain indentifier. If one exists, it is returned. If one doesn't exist,
// it is created, memory is allocated and it is returned.
func (e *Entry) getOrMakeChain(ident byte) *Chain {
	if ident == ' ' {
		ident = '_'
	}

	chain := e.Chain(ident)
	if chain != nil {
		return chain
	}
	newChain := &Chain{
		Entry:            e,
		Ident:            ident,
		Sequence:         make([]byte, 0, 30),
		AtomResidueStart: 0,
		AtomResidueEnd:   0,
		CaAtoms:          make(Atoms, 0, 30),
		CaSequence:       make([]byte, 0, 30),
	}
	e.Chains = append(e.Chains, newChain)
	return newChain
}

// parseHeader loads the "idCode" and "classification" fields from the
// header record.
//
// If the fields are already filled, then we've seen a second header record
// and therefore report an error.
func (e *Entry) parseHeader(line []byte) error {
	if len(e.Classification) > 0 || len(e.IdCode) > 0 {
		return fmt.Errorf("More than one HEADER record was found.")
	}
	e.Classification = strings.TrimSpace(string(line[10:50]))
	e.IdCode = strings.TrimSpace(string(line[62:66]))
	return nil
}

// parseSeqres loads all pertinent information from SEQRES records in a PDB
// file. In particular, amino acid resides are read and added to the chain's
// "Sequence" field. If a residue isn't a valid amino acid, it is simply
// ignored.
//
// N.B. This assumes that the SEQRES records are in order in the PDB file.
func (e *Entry) parseSeqres(line []byte) {
	chain := e.getOrMakeChain(line[11])

	// Residues are in columns 19-21, 23-25, 27-29, ..., 67-69
	for i := 19; i <= 67; i += 4 {
		end := i + 3

		// If we're passed the end of this line, quit.
		if end >= len(line) {
			break
		}

		// Get the residue. If it's not in our sequence map, skip it.
		residue := strings.TrimSpace(string(line[i:end]))
		if single, ok := AminoThreeToOne[residue]; ok {
			chain.Sequence = append(chain.Sequence, single)
		}
	}
}

// parseAtom loads all pertinent information from ATOM records in a PDB file.
// Currently, this only includes deducing the amino acid residue start and
// stop indices. (Note that the length of the range is not necessarily
// equivalent to the length of the amino acid sequence found in the SEQRES
// records.)
//
// ATOM records without a valid amino acid residue in columns 18-20 are ignored.
func (e *Entry) parseAtom(line []byte) {
	chain := e.getOrMakeChain(line[21])

	// An ATOM record is only processed if it corresponds to an amino acid
	// residue. (Which is in columns 17-19.)
	residue := strings.TrimSpace(string(line[17:20]))
	if _, ok := AminoThreeToOne[residue]; !ok {
		// Sanity check. I'm pretty sure that only amino acids have three
		// letter abbreviations.
		if len(residue) == 3 {
			panic(fmt.Sprintf("The residue '%s' found in PDB file '%s' has "+
				"length 3, but is not in my amino acid map.",
				residue, e.Path))
		}
		return
	}

	// The residue sequence number is in columns 22-25. Grab it, trim it,
	// and look for an integer.
	snum := strings.TrimSpace(string(line[22:26]))
	inum := int(0)
	if num, err := strconv.ParseInt(snum, 10, 32); err == nil {
		inum = int(num)
		switch {
		case chain.AtomResidueStart == 0 || inum < chain.AtomResidueStart:
			chain.AtomResidueStart = inum
		case chain.AtomResidueEnd == 0 || inum > chain.AtomResidueEnd:
			chain.AtomResidueEnd = inum
		}
	}

	// Build an Atom value. We need the serial number from columns 6-10,
	// the atom name from columns 12-15, the amino acid residue from
	// columns 17-19 (we already have that: 'residue'), the residue sequence
	// number from columns 22-25 (already have that too: 'inum'), and the
	// three dimension coordinates in columns 30-37 (x), 38-45 (y), and
	// 46-53 (z).
	atom := Atom{
		Name:       strings.TrimSpace(string(line[12:16])),
		Residue:    residue,
		ResidueInd: inum,
		Coords:     [3]float64{},
	}

	serialStr := strings.TrimSpace(string(line[6:11]))
	if serial64, err := strconv.ParseInt(serialStr, 10, 32); err == nil {
		atom.Serial = int(serial64)
	}

	xstr := strings.TrimSpace(string(line[30:38]))
	ystr := strings.TrimSpace(string(line[38:46]))
	zstr := strings.TrimSpace(string(line[46:54]))
	if x64, err := strconv.ParseFloat(xstr, 64); err == nil {
		atom.Coords[0] = x64
	}
	if y64, err := strconv.ParseFloat(ystr, 64); err == nil {
		atom.Coords[1] = y64
	}
	if z64, err := strconv.ParseFloat(zstr, 64); err == nil {
		atom.Coords[2] = z64
	}

	// Now add our atom to the chain.
	chain.Atoms = append(chain.Atoms, atom)
	if atom.Name == "CA" {
		chain.CaAtoms = append(chain.CaAtoms, atom)
		chain.CaSequence = append(chain.CaSequence, AminoThreeToOne[residue])
	}
}

// Chain represents a protein chain or subunit in a PDB file. Each chain has
// its own identifier, amino acid sequence (if its a protein sequence), and
// the start and stop residue indices of the ATOM coordinates.
//
// It also contains a slice of all carbon-alpha ATOM records corresponding
// to an amino acid.
type Chain struct {
	Entry                            *Entry
	Ident                            byte
	Sequence                         []byte
	AtomResidueStart, AtomResidueEnd int
	Atoms                            Atoms
	CaAtoms                          Atoms
	CaSequence                       []byte
}

// ValidProtein returns true when there are ATOM records corresponding to
// a protein backbone.
func (c *Chain) ValidProtein() bool {
	return c.AtomResidueStart != c.AtomResidueEnd
}

// String returns a FASTA-like formatted string of this chain and all of its
// related information.
func (c *Chain) String() string {
	return strings.TrimSpace(
		fmt.Sprintf("> Chain %c (%d, %d) :: length %d\n%s",
			c.Ident, c.AtomResidueStart, c.AtomResidueEnd,
			len(c.Sequence), string(c.Sequence)))
}

// Atom contains information about an ATOM record, including the serial
// number, the residue (and residue sequence number), the atom name and the
// three dimensional coordinates.
type Atom struct {
	Serial     int
	Name       string
	ResidueInd int
	Residue    string

	// Coords is a triple where the first element is X, the second is Y and
	// the third is Z.
	Coords [3]float64
}

func (a Atom) String() string {
	return fmt.Sprintf("(%d, %s, %d, %s, [%0.4f %0.4f %0.4f])",
		a.Serial, a.Name, a.ResidueInd, a.Residue,
		a.Coords[0], a.Coords[1], a.Coords[2])
}

// Atoms names a slice of Atom for sorting.
type Atoms []Atom

func (as Atoms) String() string {
	lines := make([]string, len(as))
	for i, atom := range as {
		lines[i] = atom.String()
	}
	return strings.Join(lines, "\n")
}
<commit_msg>Update pdb package to use new seq types. Update pdb2fasta to work with new types too.<commit_after>package pdb

import (
	"bufio"
	"compress/gzip"
	"fmt"
	"io"
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/BurntSushi/bcbgo/seq"
)

// AminoThreeToOne is a map from three letter amino acids to their
// corresponding single letter representation.
var AminoThreeToOne = map[string]byte{
	"ALA": 'A', "ARG": 'R', "ASN": 'N', "ASP": 'D', "CYS": 'C',
	"GLU": 'E', "GLN": 'Q', "GLY": 'G', "HIS": 'H', "ILE": 'I',
	"LEU": 'L', "LYS": 'K', "MET": 'M', "PHE": 'F', "PRO": 'P',
	"SER": 'S', "THR": 'T', "TRP": 'W', "TYR": 'Y', "VAL": 'V',
	"SEC": 'U', "PYL": 'O',
	"UNK": 'X', "ACE": 'X', "NH2": 'X',
	"ASX": 'X', "GLX": 'X',
}

// AminoOneToThree is the reverse of AminoThreeToOne. It is created in
// this packages 'init' function.
var AminoOneToThree = map[byte]string{}

func init() {
	// Create a reverse map of AminoThreeToOne.
	for k, v := range AminoThreeToOne {
		AminoOneToThree[v] = k
	}
}

// Entry represents all information known about a particular PDB file (that
// has been implemented in this package).
//
// Currently, a PDB entry is simply a file path and a map of protein chains.
type Entry struct {
	Path           string
	IdCode         string
	Classification string
	Chains         []*Chain
}

// New creates a new PDB Entry from a file. If the file cannot be read, or there
// is an error parsing the PDB file, an error is returned.
//
// If the file name ends with ".gz", gzip decompression will be used.
func New(fileName string) (*Entry, error) {
	var reader io.Reader
	var err error

	reader, err = os.Open(fileName)
	if err != nil {
		return nil, err
	}

	// If the file is gzipped, use the gzip decompressor.
	if path.Ext(fileName) == ".gz" {
		reader, err = gzip.NewReader(reader)
		if err != nil {
			return nil, err
		}
	}

	entry := &Entry{
		Path:   fileName,
		Chains: make([]*Chain, 0),
	}

	// Now traverse each line, and process it according to the record name.
	// Note that it is imperative that we preserve the order of ATOM records
	// as we read them. We are currently trying to replicate Fragbag, and this
	// is what Fragbag does. (A more stable approach would require more
	// information from the PDB file; like differentiating models, since
	// sorting on ATOM serial number isn't good enough.)
	breader := bufio.NewReaderSize(reader, 1000)
	for {
		// We ignore 'isPrefix' here, since we never care about lines longer
		// than 1000 characters, which is the size of our buffer.
		line, _, err := breader.ReadLine()
		if err == io.EOF {
			break
		} else if err != nil {
			return nil, err
		}

		// The record name is always in the fix six columns.
		switch strings.TrimSpace(string(line[0:6])) {
		case "HEADER":
			if err := entry.parseHeader(line); err != nil {
				return nil, err
			}
		case "SEQRES":
			entry.parseSeqres(line)
		case "ATOM":
			entry.parseAtom(line)
		}
	}

	// If we didn't pick up any chains, this probably isn't a valid PDB file.
	if len(entry.Chains) == 0 {
		return nil, fmt.Errorf("The file '%s' does not appear to be a valid "+
			"PDB file.", fileName)
	}

	// If we couldn't find an Id code, inspect the base name of the file path.
	if len(entry.IdCode) == 0 {
		name := path.Base(fileName)
		switch {
		case len(name) >= 7 && name[0:3] == "pdb":
			entry.IdCode = name[3:7]
		case len(name) == 7: // cath
			entry.IdCode = name[0:4]
		}
	}

	return entry, nil
}

// Chain looks for the chain with identifier ident and returns it. 'nil' is
// returned if the chain could not be found.
func (e *Entry) Chain(ident byte) *Chain {
	for _, chain := range e.Chains {
		if chain.Ident == ident {
			return chain
		}
	}
	return nil
}

// OneChain returns a single chain in the PDB file. If there is more than one
// chain, OneChain will panic. This is convenient when you expect a PDB file to
// have only a single chain, but don't know the name.
func (e *Entry) OneChain() *Chain {
	if len(e.Chains) != 1 {
		panic(fmt.Sprintf("OneChain can only be called on PDB entries with "+
			"ONE chain. But the '%s' PDB entry has %d chains.",
			e.Path, len(e.Chains)))
	}
	return e.Chains[0]
}

// Name returns the base name of the path of this PDB entry.
func (e *Entry) Name() string {
	return path.Base(e.Path)
}

// String returns a list of all chains, their residue start/stop indices,
// and the amino acid sequence.
func (e *Entry) String() string {
	lines := make([]string, 0)
	for _, chain := range e.Chains {
		lines = append(lines, chain.String())
	}
	return strings.Join(lines, "\n")
}

// getOrMakeChain looks for a chain in the 'Chains' slice corresponding to the
// chain indentifier. If one exists, it is returned. If one doesn't exist,
// it is created, memory is allocated and it is returned.
func (e *Entry) getOrMakeChain(ident byte) *Chain {
	if ident == ' ' {
		ident = '_'
	}

	chain := e.Chain(ident)
	if chain != nil {
		return chain
	}
	newChain := &Chain{
		Entry:            e,
		Ident:            ident,
		Sequence:         make([]seq.Residue, 0, 30),
		AtomResidueStart: 0,
		AtomResidueEnd:   0,
		CaAtoms:          make(Atoms, 0, 30),
		CaSequence:       make([]seq.Residue, 0, 30),
	}
	e.Chains = append(e.Chains, newChain)
	return newChain
}

// parseHeader loads the "idCode" and "classification" fields from the
// header record.
//
// If the fields are already filled, then we've seen a second header record
// and therefore report an error.
func (e *Entry) parseHeader(line []byte) error {
	if len(e.Classification) > 0 || len(e.IdCode) > 0 {
		return fmt.Errorf("More than one HEADER record was found.")
	}
	e.Classification = strings.TrimSpace(string(line[10:50]))
	e.IdCode = strings.TrimSpace(string(line[62:66]))
	return nil
}

// parseSeqres loads all pertinent information from SEQRES records in a PDB
// file. In particular, amino acid resides are read and added to the chain's
// "Sequence" field. If a residue isn't a valid amino acid, it is simply
// ignored.
//
// N.B. This assumes that the SEQRES records are in order in the PDB file.
func (e *Entry) parseSeqres(line []byte) {
	chain := e.getOrMakeChain(line[11])

	// Residues are in columns 19-21, 23-25, 27-29, ..., 67-69
	for i := 19; i <= 67; i += 4 {
		end := i + 3

		// If we're passed the end of this line, quit.
		if end >= len(line) {
			break
		}

		// Get the residue. If it's not in our sequence map, skip it.
		residue := strings.TrimSpace(string(line[i:end]))
		if single, ok := AminoThreeToOne[residue]; ok {
			chain.Sequence = append(chain.Sequence, seq.Residue(single))
		}
	}
}

// parseAtom loads all pertinent information from ATOM records in a PDB file.
// Currently, this only includes deducing the amino acid residue start and
// stop indices. (Note that the length of the range is not necessarily
// equivalent to the length of the amino acid sequence found in the SEQRES
// records.)
//
// ATOM records without a valid amino acid residue in columns 18-20 are ignored.
func (e *Entry) parseAtom(line []byte) {
	chain := e.getOrMakeChain(line[21])

	// An ATOM record is only processed if it corresponds to an amino acid
	// residue. (Which is in columns 17-19.)
	residue := strings.TrimSpace(string(line[17:20]))
	if _, ok := AminoThreeToOne[residue]; !ok {
		// Sanity check. I'm pretty sure that only amino acids have three
		// letter abbreviations.
		if len(residue) == 3 {
			panic(fmt.Sprintf("The residue '%s' found in PDB file '%s' has "+
				"length 3, but is not in my amino acid map.",
				residue, e.Path))
		}
		return
	}

	// The residue sequence number is in columns 22-25. Grab it, trim it,
	// and look for an integer.
	snum := strings.TrimSpace(string(line[22:26]))
	inum := int(0)
	if num, err := strconv.ParseInt(snum, 10, 32); err == nil {
		inum = int(num)
		switch {
		case chain.AtomResidueStart == 0 || inum < chain.AtomResidueStart:
			chain.AtomResidueStart = inum
		case chain.AtomResidueEnd == 0 || inum > chain.AtomResidueEnd:
			chain.AtomResidueEnd = inum
		}
	}

	// Build an Atom value. We need the serial number from columns 6-10,
	// the atom name from columns 12-15, the amino acid residue from
	// columns 17-19 (we already have that: 'residue'), the residue sequence
	// number from columns 22-25 (already have that too: 'inum'), and the
	// three dimension coordinates in columns 30-37 (x), 38-45 (y), and
	// 46-53 (z).
	atom := Atom{
		Name:       strings.TrimSpace(string(line[12:16])),
		Residue:    residue,
		ResidueInd: inum,
		Coords:     [3]float64{},
	}

	serialStr := strings.TrimSpace(string(line[6:11]))
	if serial64, err := strconv.ParseInt(serialStr, 10, 32); err == nil {
		atom.Serial = int(serial64)
	}

	xstr := strings.TrimSpace(string(line[30:38]))
	ystr := strings.TrimSpace(string(line[38:46]))
	zstr := strings.TrimSpace(string(line[46:54]))
	if x64, err := strconv.ParseFloat(xstr, 64); err == nil {
		atom.Coords[0] = x64
	}
	if y64, err := strconv.ParseFloat(ystr, 64); err == nil {
		atom.Coords[1] = y64
	}
	if z64, err := strconv.ParseFloat(zstr, 64); err == nil {
		atom.Coords[2] = z64
	}

	// Now add our atom to the chain.
	chain.Atoms = append(chain.Atoms, atom)
	if atom.Name == "CA" {
		chain.CaAtoms = append(chain.CaAtoms, atom)
		chain.CaSequence = append(chain.CaSequence,
			seq.Residue(AminoThreeToOne[residue]))
	}
}

// Chain represents a protein chain or subunit in a PDB file. Each chain has
// its own identifier, amino acid sequence (if its a protein sequence), and
// the start and stop residue indices of the ATOM coordinates.
//
// It also contains a slice of all carbon-alpha ATOM records corresponding
// to an amino acid.
type Chain struct {
	Entry                            *Entry
	Ident                            byte
	Sequence                         []seq.Residue
	AtomResidueStart, AtomResidueEnd int
	Atoms                            Atoms
	CaAtoms                          Atoms
	CaSequence                       []seq.Residue
}

// ValidProtein returns true when there are ATOM records corresponding to
// a protein backbone.
func (c *Chain) ValidProtein() bool {
	return c.AtomResidueStart != c.AtomResidueEnd
}

// String returns a FASTA-like formatted string of this chain and all of its
// related information.
func (c *Chain) String() string {
	return strings.TrimSpace(
		fmt.Sprintf("> Chain %c (%d, %d) :: length %d\n%s",
			c.Ident, c.AtomResidueStart, c.AtomResidueEnd,
			len(c.Sequence), c.Sequence))
}

// Atom contains information about an ATOM record, including the serial
// number, the residue (and residue sequence number), the atom name and the
// three dimensional coordinates.
type Atom struct {
	Serial     int
	Name       string
	ResidueInd int
	Residue    string

	// Coords is a triple where the first element is X, the second is Y and
	// the third is Z.
	Coords [3]float64
}

func (a Atom) String() string {
	return fmt.Sprintf("(%d, %s, %d, %s, [%0.4f %0.4f %0.4f])",
		a.Serial, a.Name, a.ResidueInd, a.Residue,
		a.Coords[0], a.Coords[1], a.Coords[2])
}

// Atoms names a slice of Atom for sorting.
type Atoms []Atom

func (as Atoms) String() string {
	lines := make([]string, len(as))
	for i, atom := range as {
		lines[i] = atom.String()
	}
	return strings.Join(lines, "\n")
}
<|endoftext|>