<commit_before>package sessions

import (
	"time"

	"gopkg.in/jackc/pgx.v2"

	"crypto/sha256"

	"fmt"

	"github.com/Unaxiom/ulogger"
	"github.com/apratheek/schemamagic"
	"github.com/twinj/uuid"
)

var globalDatabaseName string

func init() {

}

// Init accepts number of seconds after which the session expires (if set to 0, default value of 86400 is used), the timezone (if set to "", default timezone - UTC is used), and sets up the required variables. It also starts the expiry timers of existing sessions.
func Init(sessionExpiresInSecs int64, sessionLocalTimezoneName string, applicationName string, orgName string, production bool, dbHost string, dbPort uint16, databaseName string, dbUser string, dbPassword string) (*Session, error) {
	sessionObject := new(Session)
	log = ulogger.New()
	log.SetLogLevel(ulogger.InfoLevel)
	log.RemoteAvailable = production
	log.ApplicationName = applicationName + " sessions"
	log.OrganizationName = orgName
	if production {
		ulogger.RemoteURL = "https://logs.unaxiom.com/newlog"
	}
	var err error
	globalDatabaseName = databaseName
	sessionObject.db, err = schemamagic.SetupDB(dbHost, dbPort, databaseName, dbUser, dbPassword)
	if err != nil {
		return nil, err
	}

	if sessionExpiresInSecs != 0 {
		sessionExpiryTime = sessionExpiresInSecs
	}
	if sessionLocalTimezoneName != "" {
		sessionTimezoneName = sessionLocalTimezoneName
	}

	// Load the time zone here
	timezoneLocation, err = time.LoadLocation(sessionTimezoneName)
	if err != nil {
		log.Errorln("Couldn't load timezone location --> ", sessionTimezoneName, ". Error is ", err, ".")
		timezoneLocation, _ = time.LoadLocation("UTC")
	}

	// Fetch all the active sessions here
	allSessions, err := sessionObject.FetchAllSessions()
	if err != nil {
		log.Errorln("While trying to fetch all sessions during Setup, error is ", err)
		return nil, err
	}

	// Calculate the expiry timers and delete the sessions accordingly
	for _, session := range allSessions {
		go func(session SessionData) {
			<-time.After(time.Second * time.Duration(session.ExpiryIn))
			go sessionObject.DeleteSession(session.Token)
		}(session)
	}
	return sessionObject, nil
}

// SetupTable accepts all the parameters and creates/updates the Sessions table. Don't call it along with Init(). Init() and SetupTable() should be called separately.
func SetupTable(applicationName string, orgName string, production bool, dbHost string, dbPort uint16, databaseName string, dbUser string, dbPassword string, defaultSchema string) error {
	log = ulogger.New()
	log.SetLogLevel(ulogger.InfoLevel)
	log.RemoteAvailable = production
	log.ApplicationName = applicationName + " sessions"
	log.OrganizationName = orgName
	if production {
		ulogger.RemoteURL = "https://logs.unaxiom.com/newlog"
	}
	var err error
	globalDatabaseName = databaseName
	db, err = schemamagic.SetupDB(dbHost, dbPort, databaseName, dbUser, dbPassword)
	if err != nil {
		return err
	}
	err = createSessionsTable(db, defaultSchema)
	if err != nil {
		return err
	}
	return nil
}

// createSessionsTable accepts the schema of the database and creates the Sessions table using the database parameters passed during Setup. Setup needs to be called prior. Otherwise, method may crash.
func createSessionsTable(db *pgx.ConnPool, defaultSchema string) error {
	tx, err := db.Begin()
	if err != nil {
		return err
	}
	table := tableSessions(tx, defaultSchema, globalDatabaseName)
	table.Begin()
	commitErr := tx.Commit()
	if commitErr != nil {
		tx.Rollback()
		return commitErr
	}
	return nil
}

// FetchAllSessions returns all the live sessions.
func (sessionObject *Session) FetchAllSessions() ([]SessionData, error) {
	allSessions := make([]SessionData, 0)
	rows, err := sessionObject.db.Query(`
		SELECT id, key, token, expires_at, ip, timestamp FROM sessions WHERE active = True
	`)
	defer rows.Close()
	if err != nil {
		log.Errorln("While trying to fetch all sessions in FetchAllSessions, error is ", err)
		return allSessions, err
	}
	for rows.Next() {
		var session SessionData
		rows.Scan(&session.ID, &session.Key, &session.Token, &session.ExpiresAt, &session.IP, &session.Timestamp)
		session.ExpiryIn = calculateExpiresIn(session.ExpiresAt)
		allSessions = append(allSessions, session)
	}
	return allSessions, nil
}

// FetchSessionData returns the data of the particular session token
func (sessionObject *Session) FetchSessionData(session SessionData) (SessionData, error) {
	err := sessionObject.db.QueryRow(`
		SELECT id, key, token, expires_at, ip, active, timestamp FROM sessions WHERE token = $1
	`, session.Token).Scan(&session.ID, &session.Key, &session.Token, &session.ExpiresAt, &session.IP, &session.Active, &session.Timestamp)
	if err != nil {
		log.Errorln("While fetching individual session data with token --> ", session.Token, " error is ", err)
		return session, err
	}
	session.ExpiryIn = calculateExpiresIn(session.ExpiresAt)
	return session, nil
}

// NewSession accepts the key to encode, along with the client's IP Address, inserts into the sessions table, sets the delete session timer, and returns a SessionData struct.
func (sessionObject *Session) NewSession(key string, ipAddress string) (SessionData, error) {
	var assignedSession SessionData
	assignedSession.Key = key
	assignedSession.IP = ipAddress

	// Calculate the session token here
	assignedSession.Token = calculateHash(assignedSession.Key, assignedSession.IP)
	// Calculate the ExpiresAt token here
	assignedSession.ExpiresAt = time.Now().In(timezoneLocation).Add(time.Second * time.Duration(sessionExpiryTime))

	// Insert this into the database here
	tx, err := sessionObject.db.Begin()
	defer tx.Rollback()
	if err != nil {
		log.Errorln("Couldn't fetch a database connection. Error is --> ", err)
		return assignedSession, err
	}

	err = tx.QueryRow(`
		INSERT INTO sessions (key, token, expires_at, ip) VALUES ($1, $2, $3, $4) RETURNING id, timestamp
	`, assignedSession.Key, assignedSession.Token, assignedSession.ExpiresAt, assignedSession.IP).Scan(&assignedSession.ID, &assignedSession.Timestamp)

	if err != nil {
		log.Errorln("While inserting a new session in the sessions table, error is ", err)
		return assignedSession, err
	}

	commitErr := tx.Commit()
	if commitErr != nil {
		tx.Rollback()
		log.Errorln("While inserting into sessions table, error is ", commitErr)
		return assignedSession, commitErr
	}
	// Calculate SessionData.ExpiresIn from SessionData.ExpiresAt (after adjusting it to the specified time zone)
	assignedSession.ExpiryIn = calculateExpiresIn(assignedSession.ExpiresAt)

	// Set up the delete timer here
	go func() {
		<-time.After(time.Second * time.Duration(assignedSession.ExpiryIn))
		go sessionObject.DeleteSession(assignedSession.Token)
	}()

	return assignedSession, nil
}

// CheckStatus accepts a session token, and returns the session object, along with an error, if any
func (sessionObject *Session) CheckStatus(sessionToken string) (SessionData, error) {
	var session SessionData
	err := sessionObject.db.QueryRow(`
		SELECT id, key, token, expires_at, ip, timestamp FROM sessions WHERE token = $1 AND active = True
	`, sessionToken).Scan(&session.ID, &session.Key, &session.Token, &session.ExpiresAt, &session.IP, &session.Timestamp)
	if err != nil {
		return session, err
	}
	session.ExpiryIn = calculateExpiresIn(session.ExpiresAt)
	return session, nil
}

// DeleteSession deletes the session with the specific sessionToken and returns an error, if any.
func (sessionObject *Session) DeleteSession(sessionToken string) error {
	tx, _ := sessionObject.db.Begin()
	defer tx.Rollback()

	_, err := tx.Exec(`
		UPDATE sessions SET active = False WHERE token = $1
	`, sessionToken)
	if err != nil {
		tx.Rollback()
		log.Errorln("While deleting a session, error is ", err)
		return err
	}

	commitErr := tx.Commit()
	if commitErr != nil {
		tx.Rollback()
		log.Errorln("While deleting a session, error is ", commitErr)
		return commitErr
	}
	return nil
}

// calculateExpiresIn calculates the time after which a session needs to be deleted, from its full-timezone timestamp.
func calculateExpiresIn(expiresAt time.Time) int64 {
	// var expiresIn int64
	// Calculate the current time in the specified timezone
	// Find the difference between expiredAt and the current time
	// difference := expiresAt.Sub(time.Now().In(timezoneLocation))
	// Find the difference in seconds and return the value
	// return expiresIn
	// log.Infoln("Expires At timezone is ", expiresAt.Location().String(), " and timezoneLocation is ", timezoneLocation.String())
	// zone, offset := expiresAt.Zone()
	// log.Infoln("Zone is ", zone, " and offset is ", offset)
	// return int64(expiresAt.Sub(time.Now().In(timezoneLocation))) / (1000 * 1000 * 1000)
	return int64(expiresAt.Sub(time.Now())) / (1000 * 1000 * 1000)
}

// calculateHash accepts the key, and the ip address of the client, and generates an SHA
func calculateHash(key string, ipAddress string) string {
	h := sha256.New()
	h.Write([]byte(key + uuid.NewV4().String() + ipAddress))
	return fmt.Sprintf("%x", h.Sum(nil))
}
<commit_msg>If timezone location load fails, an info log is displayed, instead of an error log.<commit_after>package sessions

import (
	"time"

	"gopkg.in/jackc/pgx.v2"

	"crypto/sha256"

	"fmt"

	"github.com/Unaxiom/ulogger"
	"github.com/apratheek/schemamagic"
	"github.com/twinj/uuid"
)

var globalDatabaseName string

func init() {

}

// Init accepts number of seconds after which the session expires (if set to 0, default value of 86400 is used), the timezone (if set to "", default timezone - UTC is used), and sets up the required variables. It also starts the expiry timers of existing sessions.
func Init(sessionExpiresInSecs int64, sessionLocalTimezoneName string, applicationName string, orgName string, production bool, dbHost string, dbPort uint16, databaseName string, dbUser string, dbPassword string) (*Session, error) {
	sessionObject := new(Session)
	log = ulogger.New()
	log.SetLogLevel(ulogger.InfoLevel)
	log.RemoteAvailable = production
	log.ApplicationName = applicationName + " sessions"
	log.OrganizationName = orgName
	if production {
		ulogger.RemoteURL = "https://logs.unaxiom.com/newlog"
	}
	var err error
	globalDatabaseName = databaseName
	sessionObject.db, err = schemamagic.SetupDB(dbHost, dbPort, databaseName, dbUser, dbPassword)
	if err != nil {
		return nil, err
	}

	if sessionExpiresInSecs != 0 {
		sessionExpiryTime = sessionExpiresInSecs
	}
	if sessionLocalTimezoneName != "" {
		sessionTimezoneName = sessionLocalTimezoneName
	}

	// Load the time zone here
	timezoneLocation, err = time.LoadLocation(sessionTimezoneName)
	if err != nil {
		log.Infoln("Couldn't load timezone location --> ", sessionTimezoneName, ". Error is ", err, ".")
		timezoneLocation, _ = time.LoadLocation("UTC")
	}

	// Fetch all the active sessions here
	allSessions, err := sessionObject.FetchAllSessions()
	if err != nil {
		log.Errorln("While trying to fetch all sessions during Setup, error is ", err)
		return nil, err
	}

	// Calculate the expiry timers and delete the sessions accordingly
	for _, session := range allSessions {
		go func(session SessionData) {
			<-time.After(time.Second * time.Duration(session.ExpiryIn))
			go sessionObject.DeleteSession(session.Token)
		}(session)
	}
	return sessionObject, nil
}

// SetupTable accepts all the parameters and creates/updates the Sessions table. Don't call it along with Init(). Init() and SetupTable() should be called separately.
func SetupTable(applicationName string, orgName string, production bool, dbHost string, dbPort uint16, databaseName string, dbUser string, dbPassword string, defaultSchema string) error {
	log = ulogger.New()
	log.SetLogLevel(ulogger.InfoLevel)
	log.RemoteAvailable = production
	log.ApplicationName = applicationName + " sessions"
	log.OrganizationName = orgName
	if production {
		ulogger.RemoteURL = "https://logs.unaxiom.com/newlog"
	}
	var err error
	globalDatabaseName = databaseName
	db, err = schemamagic.SetupDB(dbHost, dbPort, databaseName, dbUser, dbPassword)
	if err != nil {
		return err
	}
	err = createSessionsTable(db, defaultSchema)
	if err != nil {
		return err
	}
	return nil
}

// createSessionsTable accepts the schema of the database and creates the Sessions table using the database parameters passed during Setup. Setup needs to be called prior. Otherwise, method may crash.
func createSessionsTable(db *pgx.ConnPool, defaultSchema string) error {
	tx, err := db.Begin()
	if err != nil {
		return err
	}
	table := tableSessions(tx, defaultSchema, globalDatabaseName)
	table.Begin()
	commitErr := tx.Commit()
	if commitErr != nil {
		tx.Rollback()
		return commitErr
	}
	return nil
}

// FetchAllSessions returns all the live sessions.
func (sessionObject *Session) FetchAllSessions() ([]SessionData, error) {
	allSessions := make([]SessionData, 0)
	rows, err := sessionObject.db.Query(`
		SELECT id, key, token, expires_at, ip, timestamp FROM sessions WHERE active = True
	`)
	defer rows.Close()
	if err != nil {
		log.Errorln("While trying to fetch all sessions in FetchAllSessions, error is ", err)
		return allSessions, err
	}
	for rows.Next() {
		var session SessionData
		rows.Scan(&session.ID, &session.Key, &session.Token, &session.ExpiresAt, &session.IP, &session.Timestamp)
		session.ExpiryIn = calculateExpiresIn(session.ExpiresAt)
		allSessions = append(allSessions, session)
	}
	return allSessions, nil
}

// FetchSessionData returns the data of the particular session token
func (sessionObject *Session) FetchSessionData(session SessionData) (SessionData, error) {
	err := sessionObject.db.QueryRow(`
		SELECT id, key, token, expires_at, ip, active, timestamp FROM sessions WHERE token = $1
	`, session.Token).Scan(&session.ID, &session.Key, &session.Token, &session.ExpiresAt, &session.IP, &session.Active, &session.Timestamp)
	if err != nil {
		log.Errorln("While fetching individual session data with token --> ", session.Token, " error is ", err)
		return session, err
	}
	session.ExpiryIn = calculateExpiresIn(session.ExpiresAt)
	return session, nil
}

// NewSession accepts the key to encode, along with the client's IP Address, inserts into the sessions table, sets the delete session timer, and returns a SessionData struct.
func (sessionObject *Session) NewSession(key string, ipAddress string) (SessionData, error) {
	var assignedSession SessionData
	assignedSession.Key = key
	assignedSession.IP = ipAddress

	// Calculate the session token here
	assignedSession.Token = calculateHash(assignedSession.Key, assignedSession.IP)
	// Calculate the ExpiresAt token here
	assignedSession.ExpiresAt = time.Now().In(timezoneLocation).Add(time.Second * time.Duration(sessionExpiryTime))

	// Insert this into the database here
	tx, err := sessionObject.db.Begin()
	defer tx.Rollback()
	if err != nil {
		log.Errorln("Couldn't fetch a database connection. Error is --> ", err)
		return assignedSession, err
	}

	err = tx.QueryRow(`
		INSERT INTO sessions (key, token, expires_at, ip) VALUES ($1, $2, $3, $4) RETURNING id, timestamp
	`, assignedSession.Key, assignedSession.Token, assignedSession.ExpiresAt, assignedSession.IP).Scan(&assignedSession.ID, &assignedSession.Timestamp)

	if err != nil {
		log.Errorln("While inserting a new session in the sessions table, error is ", err)
		return assignedSession, err
	}

	commitErr := tx.Commit()
	if commitErr != nil {
		tx.Rollback()
		log.Errorln("While inserting into sessions table, error is ", commitErr)
		return assignedSession, commitErr
	}
	// Calculate SessionData.ExpiresIn from SessionData.ExpiresAt (after adjusting it to the specified time zone)
	assignedSession.ExpiryIn = calculateExpiresIn(assignedSession.ExpiresAt)

	// Set up the delete timer here
	go func() {
		<-time.After(time.Second * time.Duration(assignedSession.ExpiryIn))
		go sessionObject.DeleteSession(assignedSession.Token)
	}()

	return assignedSession, nil
}

// CheckStatus accepts a session token, and returns the session object, along with an error, if any
func (sessionObject *Session) CheckStatus(sessionToken string) (SessionData, error) {
	var session SessionData
	err := sessionObject.db.QueryRow(`
		SELECT id, key, token, expires_at, ip, timestamp FROM sessions WHERE token = $1 AND active = True
	`, sessionToken).Scan(&session.ID, &session.Key, &session.Token, &session.ExpiresAt, &session.IP, &session.Timestamp)
	if err != nil {
		return session, err
	}
	session.ExpiryIn = calculateExpiresIn(session.ExpiresAt)
	return session, nil
}

// DeleteSession deletes the session with the specific sessionToken and returns an error, if any.
func (sessionObject *Session) DeleteSession(sessionToken string) error {
	tx, _ := sessionObject.db.Begin()
	defer tx.Rollback()

	_, err := tx.Exec(`
		UPDATE sessions SET active = False WHERE token = $1
	`, sessionToken)
	if err != nil {
		tx.Rollback()
		log.Errorln("While deleting a session, error is ", err)
		return err
	}

	commitErr := tx.Commit()
	if commitErr != nil {
		tx.Rollback()
		log.Errorln("While deleting a session, error is ", commitErr)
		return commitErr
	}
	return nil
}

// calculateExpiresIn calculates the time after which a session needs to be deleted, from its full-timezone timestamp.
func calculateExpiresIn(expiresAt time.Time) int64 {
	// var expiresIn int64
	// Calculate the current time in the specified timezone
	// Find the difference between expiredAt and the current time
	// difference := expiresAt.Sub(time.Now().In(timezoneLocation))
	// Find the difference in seconds and return the value
	// return expiresIn
	// log.Infoln("Expires At timezone is ", expiresAt.Location().String(), " and timezoneLocation is ", timezoneLocation.String())
	// zone, offset := expiresAt.Zone()
	// log.Infoln("Zone is ", zone, " and offset is ", offset)
	// return int64(expiresAt.Sub(time.Now().In(timezoneLocation))) / (1000 * 1000 * 1000)
	return int64(expiresAt.Sub(time.Now())) / (1000 * 1000 * 1000)
}

// calculateHash accepts the key, and the ip address of the client, and generates an SHA
func calculateHash(key string, ipAddress string) string {
	h := sha256.New()
	h.Write([]byte(key + uuid.NewV4().String() + ipAddress))
	return fmt.Sprintf("%x", h.Sum(nil))
}
<|endoftext|>


<commit_before>// Package sessions contains middleware for easy session management in Martini.
//
//  package main
//
//  import (
//    "github.com/codegangsta/martini"
//    "github.com/codegangsta/martini-contrib/sessions"
//  )
//
//  func main() {
// 	  m := martini.Classic()
//
// 	  store := sessions.NewCookieStore([]byte("secret123"))
// 	  m.Use(sessions.Sessions("my_session", store))
//
// 	  m.Get("/", func(session sessions.Session) string {
// 		  session.Set("hello", "world")
// 	  })
//  }
package sessions

import (
	"github.com/codegangsta/martini"
	"github.com/gorilla/context"
	"github.com/gorilla/sessions"
	"log"
	"net/http"
)

const (
	errorFormat = "[sessions] ERROR! %s\n"
)

// Store is an interface for custom session stores.
type Store interface {
	sessions.Store
}

// NewCookieStore returns a new CookieStore.
//
// Keys are defined in pairs to allow key rotation, but the common case is to set a single
// authentication key and optionally an encryption key.
//
// The first key in a pair is used for authentication and the second for encryption. The
// encryption key can be set to nil or omitted in the last pair, but the authentication key
// is required in all pairs.
//
// It is recommended to use an authentication key with 32 or 64 bytes. The encryption key,
// if set, must be either 16, 24, or 32 bytes to select AES-128, AES-192, or AES-256 modes.
func NewCookieStore(keyPairs ...[]byte) Store {
	return sessions.NewCookieStore(keyPairs...)
}

// Session stores the values and optional configuration for a session.
type Session interface {
	// Get returns the session value associated to the given key.
	Get(key interface{}) interface{}
	// Set sets the session value associated to the given key.
	Set(key interface{}, val interface{})
	// AddFlash adds a flash message to the session.
	// A single variadic argument is accepted, and it is optional: it defines the flash key.
	// If not defined "_flash" is used by default.
	//AddFlash(value interface{}, vars ...string)
	// Flashes returns a slice of flash messages from the session.
	// A single variadic argument is accepted, and it is optional: it defines the flash key.
	// If not defined "_flash" is used by default.
	//Flashes(vars ...string) []interface{}
}

// Sessions is a Middleware that maps a session.Session service into the Martini handler chain.
// Sessions can use a number of storage solutions with the given store.
func Sessions(name string, store Store) martini.Handler {
	return func(res http.ResponseWriter, r *http.Request, c martini.Context, l *log.Logger) {
		// Map to the Session interface
    s := &session{name, r, l, store, nil, false}
		c.MapTo(s, (*Session)(nil))

		// clear the context, we don't need to use
		// gorilla context and we don't want memory leaks
		defer context.Clear(r)

		// Use before hook to save out the session
		rw := res.(martini.ResponseWriter)
		rw.Before(func(martini.ResponseWriter) {
      if s.Written {
        check(s.Session.Save(r, res), l)
      }
		})
	}
}

type session struct {
  Name string
  Request *http.Request
  Logger *log.Logger
  Store Store
  Session *sessions.Session
  Written bool
}

func (s *session) Get(key interface{}) interface{} {
  s.EnsureSession()
	return s.Session.Values[key]
}

func (s *session) Set(key interface{}, val interface{}) {
  s.EnsureSession()
  s.Written = true
	s.Session.Values[key] = val
}

func (s *session) EnsureSession() {
  if s.Session == nil {
    s.Session, _ = s.Store.Get(s.Request, s.Name)
  }
}

func check(err error, l *log.Logger) {
	if err != nil {
		l.Printf(errorFormat, err)
	}
}
<commit_msg>-n [Migrated] JMS #38: Formatting<commit_after>// Package sessions contains middleware for easy session management in Martini.
//
//  package main
//
//  import (
//    "github.com/codegangsta/martini"
//    "github.com/codegangsta/martini-contrib/sessions"
//  )
//
//  func main() {
// 	  m := martini.Classic()
//
// 	  store := sessions.NewCookieStore([]byte("secret123"))
// 	  m.Use(sessions.Sessions("my_session", store))
//
// 	  m.Get("/", func(session sessions.Session) string {
// 		  session.Set("hello", "world")
// 	  })
//  }
package sessions

import (
	"github.com/codegangsta/martini"
	"github.com/gorilla/context"
	"github.com/gorilla/sessions"
	"log"
	"net/http"
)

const (
	errorFormat = "[sessions] ERROR! %s\n"
)

// Store is an interface for custom session stores.
type Store interface {
	sessions.Store
}

// NewCookieStore returns a new CookieStore.
//
// Keys are defined in pairs to allow key rotation, but the common case is to set a single
// authentication key and optionally an encryption key.
//
// The first key in a pair is used for authentication and the second for encryption. The
// encryption key can be set to nil or omitted in the last pair, but the authentication key
// is required in all pairs.
//
// It is recommended to use an authentication key with 32 or 64 bytes. The encryption key,
// if set, must be either 16, 24, or 32 bytes to select AES-128, AES-192, or AES-256 modes.
func NewCookieStore(keyPairs ...[]byte) Store {
	return sessions.NewCookieStore(keyPairs...)
}

// Session stores the values and optional configuration for a session.
type Session interface {
	// Get returns the session value associated to the given key.
	Get(key interface{}) interface{}
	// Set sets the session value associated to the given key.
	Set(key interface{}, val interface{})
	// AddFlash adds a flash message to the session.
	// A single variadic argument is accepted, and it is optional: it defines the flash key.
	// If not defined "_flash" is used by default.
	//AddFlash(value interface{}, vars ...string)
	// Flashes returns a slice of flash messages from the session.
	// A single variadic argument is accepted, and it is optional: it defines the flash key.
	// If not defined "_flash" is used by default.
	//Flashes(vars ...string) []interface{}
}

// Sessions is a Middleware that maps a session.Session service into the Martini handler chain.
// Sessions can use a number of storage solutions with the given store.
func Sessions(name string, store Store) martini.Handler {
	return func(res http.ResponseWriter, r *http.Request, c martini.Context, l *log.Logger) {
		// Map to the Session interface
		s := &session{name, r, l, store, nil, false}
		c.MapTo(s, (*Session)(nil))

		// clear the context, we don't need to use
		// gorilla context and we don't want memory leaks
		defer context.Clear(r)

		// Use before hook to save out the session
		rw := res.(martini.ResponseWriter)
		rw.Before(func(martini.ResponseWriter) {
			if s.Written {
				check(s.Session.Save(r, res), l)
			}
		})
	}
}

type session struct {
	Name    string
	Request *http.Request
	Logger  *log.Logger
	Store   Store
	Session *sessions.Session
	Written bool
}

func (s *session) Get(key interface{}) interface{} {
	s.EnsureSession()
	return s.Session.Values[key]
}

func (s *session) Set(key interface{}, val interface{}) {
	s.EnsureSession()
	s.Session.Values[key] = val
	s.Written = true
}

func (s *session) EnsureSession() {
	if s.Session == nil {
		s.Session, _ = s.Store.Get(s.Request, s.Name)
	}
}

func check(err error, l *log.Logger) {
	if err != nil {
		l.Printf(errorFormat, err)
	}
}
<|endoftext|>


<commit_before>package pop

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"time"

	"github.com/pkg/errors"
)

// MigrationCreate writes contents for a given migration in normalized files
func MigrationCreate(path, name, ext string, up, down []byte) error {
	n := time.Now().UTC()
	s := n.Format("20060102150405")

	err := os.MkdirAll(path, 0766)
	if err != nil {
		return errors.Wrapf(err, "couldn't create migrations path %s", path)
	}

	upf := filepath.Join(path, (fmt.Sprintf("%s_%s.up.%s", s, name, ext)))
	err = ioutil.WriteFile(upf, up, 0666)
	if err != nil {
		return errors.Wrapf(err, "couldn't write up migration %s", upf)
	}
	fmt.Printf("> %s\n", upf)

	downf := filepath.Join(path, (fmt.Sprintf("%s_%s.down.%s", s, name, ext)))
	err = ioutil.WriteFile(downf, down, 0666)
	if err != nil {
		return errors.Wrapf(err, "couldn't write up migration %s", downf)
	}

	fmt.Printf("> %s\n", downf)
	return nil
}

// MigrateUp is deprecated, and will be removed in a future version. Use FileMigrator#Up instead.
func (c *Connection) MigrateUp(path string) error {
	warningMsg := "Connection#MigrateUp is deprecated, and will be removed in a future version. Use FileMigrator#Up instead."
	_, file, no, ok := runtime.Caller(1)
	if ok {
		warningMsg = fmt.Sprintf("%s Called from %s:%d", warningMsg, file, no)
	}
	log.Println(warningMsg)

	mig, err := NewFileMigrator(path, c)
	if err != nil {
		return errors.WithStack(err)
	}
	return mig.Up()
}

// MigrateDown is deprecated, and will be removed in a future version. Use FileMigrator#Down instead.
func (c *Connection) MigrateDown(path string, step int) error {
	warningMsg := "Connection#MigrateDown is deprecated, and will be removed in a future version. Use FileMigrator#Down instead."
	_, file, no, ok := runtime.Caller(1)
	if ok {
		warningMsg = fmt.Sprintf("%s Called from %s:%d", warningMsg, file, no)
	}
	log.Println(warningMsg)

	mig, err := NewFileMigrator(path, c)
	if err != nil {
		return errors.WithStack(err)
	}
	return mig.Down(step)
}

// MigrateStatus is deprecated, and will be removed in a future version. Use FileMigrator#Status instead.
func (c *Connection) MigrateStatus(path string) error {
	warningMsg := "Connection#MigrateStatus is deprecated, and will be removed in a future version. Use FileMigrator#Status instead."
	_, file, no, ok := runtime.Caller(1)
	if ok {
		warningMsg = fmt.Sprintf("%s Called from %s:%d", warningMsg, file, no)
	}
	log.Println(warningMsg)

	mig, err := NewFileMigrator(path, c)
	if err != nil {
		return errors.WithStack(err)
	}
	return mig.Status()
}

// MigrateReset is deprecated, and will be removed in a future version. Use FileMigrator#Reset instead.
func (c *Connection) MigrateReset(path string) error {
	warningMsg := "Connection#MigrateReset is deprecated, and will be removed in a future version. Use FileMigrator#Reset instead."
	_, file, no, ok := runtime.Caller(1)
	if ok {
		warningMsg = fmt.Sprintf("%s Called from %s:%d", warningMsg, file, no)
	}
	log.Println(warningMsg)

	mig, err := NewFileMigrator(path, c)
	if err != nil {
		return errors.WithStack(err)
	}
	return mig.Reset()
}
<commit_msg>Replace MigrationCreate manual file creation with Makr calls<commit_after>package pop

import (
	"fmt"
	"log"
	"path/filepath"
	"runtime"
	"time"

	"github.com/gobuffalo/makr"
	"github.com/pkg/errors"
)

// MigrationCreate writes contents for a given migration in normalized files
func MigrationCreate(path, name, ext string, up, down []byte) error {
	g := makr.New()
	n := time.Now().UTC()
	s := n.Format("20060102150405")

	upf := filepath.Join(path, (fmt.Sprintf("%s_%s.up.%s", s, name, ext)))
	g.Add(makr.NewFile(upf, string(up)))

	downf := filepath.Join(path, (fmt.Sprintf("%s_%s.down.%s", s, name, ext)))
	g.Add(makr.NewFile(downf, string(down)))

	return g.Run(".", makr.Data{})
}

// MigrateUp is deprecated, and will be removed in a future version. Use FileMigrator#Up instead.
func (c *Connection) MigrateUp(path string) error {
	warningMsg := "Connection#MigrateUp is deprecated, and will be removed in a future version. Use FileMigrator#Up instead."
	_, file, no, ok := runtime.Caller(1)
	if ok {
		warningMsg = fmt.Sprintf("%s Called from %s:%d", warningMsg, file, no)
	}
	log.Println(warningMsg)

	mig, err := NewFileMigrator(path, c)
	if err != nil {
		return errors.WithStack(err)
	}
	return mig.Up()
}

// MigrateDown is deprecated, and will be removed in a future version. Use FileMigrator#Down instead.
func (c *Connection) MigrateDown(path string, step int) error {
	warningMsg := "Connection#MigrateDown is deprecated, and will be removed in a future version. Use FileMigrator#Down instead."
	_, file, no, ok := runtime.Caller(1)
	if ok {
		warningMsg = fmt.Sprintf("%s Called from %s:%d", warningMsg, file, no)
	}
	log.Println(warningMsg)

	mig, err := NewFileMigrator(path, c)
	if err != nil {
		return errors.WithStack(err)
	}
	return mig.Down(step)
}

// MigrateStatus is deprecated, and will be removed in a future version. Use FileMigrator#Status instead.
func (c *Connection) MigrateStatus(path string) error {
	warningMsg := "Connection#MigrateStatus is deprecated, and will be removed in a future version. Use FileMigrator#Status instead."
	_, file, no, ok := runtime.Caller(1)
	if ok {
		warningMsg = fmt.Sprintf("%s Called from %s:%d", warningMsg, file, no)
	}
	log.Println(warningMsg)

	mig, err := NewFileMigrator(path, c)
	if err != nil {
		return errors.WithStack(err)
	}
	return mig.Status()
}

// MigrateReset is deprecated, and will be removed in a future version. Use FileMigrator#Reset instead.
func (c *Connection) MigrateReset(path string) error {
	warningMsg := "Connection#MigrateReset is deprecated, and will be removed in a future version. Use FileMigrator#Reset instead."
	_, file, no, ok := runtime.Caller(1)
	if ok {
		warningMsg = fmt.Sprintf("%s Called from %s:%d", warningMsg, file, no)
	}
	log.Println(warningMsg)

	mig, err := NewFileMigrator(path, c)
	if err != nil {
		return errors.WithStack(err)
	}
	return mig.Reset()
}
<|endoftext|>


<commit_before>package vcd

import (
	"fmt"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
	types "github.com/ukcloud/govcloudair/types/v56"
	"log"
)

func resourceVcdEdgeGatewayVpn() *schema.Resource {
	return &schema.Resource{
		Create: resourceVcdEdgeGatewayVpnCreate,
		Read:   resourceVcdEdgeGatewayVpnRead,
		Delete: resourceVcdEdgeGatewayVpnDelete,

		Schema: map[string]*schema.Schema{

			"edge_gateway": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"description": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},

			"encryption_protocol": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"local_ip_address": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"local_id": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"mtu": &schema.Schema{
				Type:     schema.TypeInt,
				Required: true,
				ForceNew: true,
			},

			"peer_ip_address": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"peer_id": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"shared_secret": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"local_subnets": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				ForceNew: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"local_subnet_name": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"local_subnet_gateway": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"local_subnet_mask": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},
					},
				},
			},

			"peer_subnets": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				ForceNew: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"peer_subnet_name": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"peer_subnet_gateway": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"peer_subnet_mask": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},
					},
				},
			},
		},
	}
}

func resourceVcdEdgeGatewayVpnCreate(d *schema.ResourceData, meta interface{}) error {
	vcdClient := meta.(*VCDClient)
	log.Printf("[TRACE] CLIENT: %#v", vcdClient)
	vcdClient.Mutex.Lock()
	defer vcdClient.Mutex.Unlock()

	edgeGateway, err := vcdClient.OrgVdc.FindEdgeGateway(d.Get("edge_gateway").(string))

	localSubnetsList := d.Get("local_subnets").(*schema.Set).List()
	peerSubnetsList := d.Get("peer_subnets").(*schema.Set).List()

	localSubnets := make([]*types.IpsecVpnSubnet, len(localSubnetsList))
	peerSubnets := make([]*types.IpsecVpnSubnet, len(peerSubnetsList))

	for i, s := range localSubnetsList {
		ls := s.(map[string]interface{})
		localSubnets[i] = &types.IpsecVpnSubnet{
			Name:    ls["local_subnet_name"].(string),
			Gateway: ls["local_subnet_gateway"].(string),
			Netmask: ls["local_subnet_mask"].(string),
		}
	}

	for i, s := range peerSubnetsList {
		ls := s.(map[string]interface{})
		peerSubnets[i] = &types.IpsecVpnSubnet{
			Name:    ls["peer_subnet_name"].(string),
			Gateway: ls["peer_subnet_gateway"].(string),
			Netmask: ls["peer_subnet_mask"].(string),
		}
	}

	tunnel := &types.GatewayIpsecVpnTunnel{
		Name:        d.Get("name").(string),
		Description: d.Get("description").(string),
		IpsecVpnLocalPeer: &types.IpsecVpnLocalPeer{
			ID:   "",
			Name: "",
		},
		EncryptionProtocol: d.Get("encryption_protocol").(string),
		LocalIPAddress:     d.Get("local_ip_address").(string),
		LocalID:            d.Get("local_id").(string),
		LocalSubnet:        localSubnets,
		Mtu:                d.Get("mtu").(int),
		PeerID:             d.Get("peer_id").(string),
		PeerIPAddress:      d.Get("peer_ip_address").(string),
		PeerSubnet:         peerSubnets,
		SharedSecret:       d.Get("shared_secret").(string),
		IsEnabled:          true,
	}

	tunnels := make([]*types.GatewayIpsecVpnTunnel, 1)
	tunnels[0] = tunnel

	ipsecVPNConfig := &types.EdgeGatewayServiceConfiguration{
		Xmlns: "http://www.vmware.com/vcloud/v1.5",
		GatewayIpsecVpnService: &types.GatewayIpsecVpnService{
			IsEnabled: true,
			Tunnel:    tunnels,
		},
	}

	log.Printf("[INFO] ipsecVPNConfig: %#v", ipsecVPNConfig)

	err = retryCall(vcdClient.MaxRetryTimeout, func() *resource.RetryError {
		edgeGateway.Refresh()
		task, err := edgeGateway.AddIpsecVPN(ipsecVPNConfig)
		if err != nil {
			log.Printf("[INFO] Error setting ipsecVPNConfig rules: %s", err)
			return resource.RetryableError(
				fmt.Errorf("Error setting ipsecVPNConfig rules: %#v", err))
		}

		return resource.RetryableError(task.WaitTaskCompletion())
	})
	if err != nil {
		return fmt.Errorf("Error completing tasks: %#v", err)
	}

	d.SetId(d.Get("edge_gateway").(string))

	return resourceVcdEdgeGatewayVpnRead(d, meta)
}

func resourceVcdEdgeGatewayVpnDelete(d *schema.ResourceData, meta interface{}) error {
	vcdClient := meta.(*VCDClient)

	log.Printf("[TRACE] CLIENT: %#v", vcdClient)

	vcdClient.Mutex.Lock()
	defer vcdClient.Mutex.Unlock()

	edgeGateway, err := vcdClient.OrgVdc.FindEdgeGateway(d.Get("edge_gateway").(string))

	ipsecVPNConfig := &types.EdgeGatewayServiceConfiguration{
		Xmlns: "http://www.vmware.com/vcloud/v1.5",
		GatewayIpsecVpnService: &types.GatewayIpsecVpnService{
			IsEnabled: false,
		},
	}

	log.Printf("[INFO] ipsecVPNConfig: %#v", ipsecVPNConfig)

	err = retryCall(vcdClient.MaxRetryTimeout, func() *resource.RetryError {
		edgeGateway.Refresh()
		task, err := edgeGateway.AddIpsecVPN(ipsecVPNConfig)
		if err != nil {
			log.Printf("[INFO] Error setting ipsecVPNConfig rules: %s", err)
			return resource.RetryableError(
				fmt.Errorf("Error setting ipsecVPNConfig rules: %#v", err))
		}

		return resource.RetryableError(task.WaitTaskCompletion())
	})
	if err != nil {
		return fmt.Errorf("Error completing tasks: %#v", err)
	}

	d.SetId(d.Get("edge_gateway").(string))

	if err != nil {
		return fmt.Errorf("Error finding edge gateway: %#v", err)
	}

	return nil
}

func resourceVcdEdgeGatewayVpnRead(d *schema.ResourceData, meta interface{}) error {
	vcdClient := meta.(*VCDClient)

	edgeGateway, err := vcdClient.OrgVdc.FindEdgeGateway(d.Get("edge_gateway").(string))
	if err != nil {
		return fmt.Errorf("Error finding edge gateway: %#v", err)
	}

	fmt.Println(edgeGateway)

	//d.Set("name", *edgeGateway.EdgeGateway.Configuration.EdgeGatewayServiceConfiguration.GatewayIpsecVpnService.Tunnel[0].Name)
	// and all the others

	return nil
}
<commit_msg>adding read methods<commit_after>package vcd

import (
	"fmt"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
	types "github.com/ukcloud/govcloudair/types/v56"
	"log"
)

func resourceVcdEdgeGatewayVpn() *schema.Resource {
	return &schema.Resource{
		Create: resourceVcdEdgeGatewayVpnCreate,
		Read:   resourceVcdEdgeGatewayVpnRead,
		Delete: resourceVcdEdgeGatewayVpnDelete,

		Schema: map[string]*schema.Schema{

			"edge_gateway": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"description": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},

			"encryption_protocol": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"local_ip_address": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"local_id": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"mtu": &schema.Schema{
				Type:     schema.TypeInt,
				Required: true,
				ForceNew: true,
			},

			"peer_ip_address": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"peer_id": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"shared_secret": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"local_subnets": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				ForceNew: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"local_subnet_name": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"local_subnet_gateway": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"local_subnet_mask": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},
					},
				},
			},

			"peer_subnets": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				ForceNew: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"peer_subnet_name": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"peer_subnet_gateway": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"peer_subnet_mask": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},
					},
				},
			},
		},
	}
}

func resourceVcdEdgeGatewayVpnCreate(d *schema.ResourceData, meta interface{}) error {
	vcdClient := meta.(*VCDClient)
	log.Printf("[TRACE] CLIENT: %#v", vcdClient)
	vcdClient.Mutex.Lock()
	defer vcdClient.Mutex.Unlock()

	edgeGateway, err := vcdClient.OrgVdc.FindEdgeGateway(d.Get("edge_gateway").(string))

	localSubnetsList := d.Get("local_subnets").(*schema.Set).List()
	peerSubnetsList := d.Get("peer_subnets").(*schema.Set).List()

	localSubnets := make([]*types.IpsecVpnSubnet, len(localSubnetsList))
	peerSubnets := make([]*types.IpsecVpnSubnet, len(peerSubnetsList))

	for i, s := range localSubnetsList {
		ls := s.(map[string]interface{})
		localSubnets[i] = &types.IpsecVpnSubnet{
			Name:    ls["local_subnet_name"].(string),
			Gateway: ls["local_subnet_gateway"].(string),
			Netmask: ls["local_subnet_mask"].(string),
		}
	}

	for i, s := range peerSubnetsList {
		ls := s.(map[string]interface{})
		peerSubnets[i] = &types.IpsecVpnSubnet{
			Name:    ls["peer_subnet_name"].(string),
			Gateway: ls["peer_subnet_gateway"].(string),
			Netmask: ls["peer_subnet_mask"].(string),
		}
	}

	tunnel := &types.GatewayIpsecVpnTunnel{
		Name:        d.Get("name").(string),
		Description: d.Get("description").(string),
		IpsecVpnLocalPeer: &types.IpsecVpnLocalPeer{
			ID:   "",
			Name: "",
		},
		EncryptionProtocol: d.Get("encryption_protocol").(string),
		LocalIPAddress:     d.Get("local_ip_address").(string),
		LocalID:            d.Get("local_id").(string),
		LocalSubnet:        localSubnets,
		Mtu:                d.Get("mtu").(int),
		PeerID:             d.Get("peer_id").(string),
		PeerIPAddress:      d.Get("peer_ip_address").(string),
		PeerSubnet:         peerSubnets,
		SharedSecret:       d.Get("shared_secret").(string),
		IsEnabled:          true,
	}

	tunnels := make([]*types.GatewayIpsecVpnTunnel, 1)
	tunnels[0] = tunnel

	ipsecVPNConfig := &types.EdgeGatewayServiceConfiguration{
		Xmlns: "http://www.vmware.com/vcloud/v1.5",
		GatewayIpsecVpnService: &types.GatewayIpsecVpnService{
			IsEnabled: true,
			Tunnel:    tunnels,
		},
	}

	log.Printf("[INFO] ipsecVPNConfig: %#v", ipsecVPNConfig)

	err = retryCall(vcdClient.MaxRetryTimeout, func() *resource.RetryError {
		edgeGateway.Refresh()
		task, err := edgeGateway.AddIpsecVPN(ipsecVPNConfig)
		if err != nil {
			log.Printf("[INFO] Error setting ipsecVPNConfig rules: %s", err)
			return resource.RetryableError(
				fmt.Errorf("Error setting ipsecVPNConfig rules: %#v", err))
		}

		return resource.RetryableError(task.WaitTaskCompletion())
	})
	if err != nil {
		return fmt.Errorf("Error completing tasks: %#v", err)
	}

	d.SetId(d.Get("edge_gateway").(string))

	return resourceVcdEdgeGatewayVpnRead(d, meta)
}

func resourceVcdEdgeGatewayVpnDelete(d *schema.ResourceData, meta interface{}) error {
	vcdClient := meta.(*VCDClient)

	log.Printf("[TRACE] CLIENT: %#v", vcdClient)

	vcdClient.Mutex.Lock()
	defer vcdClient.Mutex.Unlock()

	edgeGateway, err := vcdClient.OrgVdc.FindEdgeGateway(d.Get("edge_gateway").(string))

	ipsecVPNConfig := &types.EdgeGatewayServiceConfiguration{
		Xmlns: "http://www.vmware.com/vcloud/v1.5",
		GatewayIpsecVpnService: &types.GatewayIpsecVpnService{
			IsEnabled: false,
		},
	}

	log.Printf("[INFO] ipsecVPNConfig: %#v", ipsecVPNConfig)

	err = retryCall(vcdClient.MaxRetryTimeout, func() *resource.RetryError {
		edgeGateway.Refresh()
		task, err := edgeGateway.AddIpsecVPN(ipsecVPNConfig)
		if err != nil {
			log.Printf("[INFO] Error setting ipsecVPNConfig rules: %s", err)
			return resource.RetryableError(
				fmt.Errorf("Error setting ipsecVPNConfig rules: %#v", err))
		}

		return resource.RetryableError(task.WaitTaskCompletion())
	})
	if err != nil {
		return fmt.Errorf("Error completing tasks: %#v", err)
	}

	d.SetId(d.Get("edge_gateway").(string))

	if err != nil {
		return fmt.Errorf("Error finding edge gateway: %#v", err)
	}

	return nil
}

func resourceVcdEdgeGatewayVpnRead(d *schema.ResourceData, meta interface{}) error {
	vcdClient := meta.(*VCDClient)

	edgeGateway, err := vcdClient.OrgVdc.FindEdgeGateway(d.Get("edge_gateway").(string))
	if err != nil {
		return fmt.Errorf("Error finding edge gateway: %#v", err)
	}

	egsc := edgeGateway.EdgeGateway.Configuration.EdgeGatewayServiceConfiguration.GatewayIpsecVpnService
	for _, t := range egsc.Tunnel {
		d.Set("name", t.Name)
		d.Set("description", t.Description)
		d.Set("encryption_protocol", t.EncryptionProtocol)
		d.Set("local_ip_address", t.LocalIPAddress)
		d.Set("local_id", t.LocalID)
		d.Set("mtu", t.Mtu)
		d.Set("peer_ip_address", t.PeerIPAddress)
		d.Set("peer_id", t.PeerID)
		d.Set("shared_secret", t.SharedSecret)
		d.Set("local_subnets", t.LocalSubnet)
		d.Set("peer_subnets", t.PeerSubnet)

	}
	// and all the others

	return nil
}
<|endoftext|>


<commit_before>// minicache is a work in progress in-memory caching system
// featuring a similar text based protocol to memcached
package main

// TODO:
// - TTL handling
// - casunique?
// - noreply?
// - Investigate how memcached handles concurrency/locking

import (
	"bufio"
	"fmt"
	"log"
	"net"
	// "os"
	"strconv"
	"strings"
	"time"
)

type Client struct {
	Id      string
	State   uint8
	Input   []string
	Command string
	Record  *Record
}

type Record struct {
	Key    string
	Value  []byte
	Flags  int64
	Ttl    int64
	Length int64
}

const STATE_DEFAULT uint8 = 1
const STATE_EXPECTING_VALUE uint8 = 2

const STATE_COMMAND_GET uint8 = 3
const STATE_COMMAND_SET uint8 = 4
const STATE_COMMAND_DELETE uint8 = 5
const STATE_COMMAND_QUIT uint8 = 6
const STATE_COMMAND_FLUSH_ALL uint8 = 7

var ticker = time.NewTicker(time.Second * 1)
var clients map[string]*Client
var datastore map[string]*Record

func main() {
	// Start the server on port 5268
	server, err := net.Listen("tcp", ":5268")

	if err != nil {
		log.Fatal(err)
	}

	// Ensure that the server closes
	defer server.Close()

	// Our datastore..
	datastore = make(map[string]*Record)

	// A list of active clients
	clients = make(map[string]*Client)

	// Print out the datastore contents every second for debug
	go func() {
		for range ticker.C {
			fmt.Println(datastore)
		}
	}()

	for {
		// Wait for a connection.
		connection, err := server.Accept()

		if err != nil {
			log.Fatal(err)
		}

		// Handle the connection in a new goroutine.
		// The loop then returns to accepting new connections,
		// so that multiple connections may be served concurrently.
		go func(connection net.Conn) {
			// Create the client and store it
			client := &Client{
				Id:      connection.RemoteAddr().String(),
				State:   STATE_DEFAULT,
				Input:   []string{},
				Command: "",
				Record:  &Record{},
			}
			clients[connection.RemoteAddr().String()] = client

			// Ensure the client is tidied up once they're done
			defer func(connection net.Conn, clients map[string]*Client, id string) {
				delete(clients, id)
				connection.Close()
			}(connection, clients, client.Id)

			// Create a new scanner for the client input
			scanner := bufio.NewScanner(connection)

			// Handle each line (command)
			for scanner.Scan() {
				// Split the client input up based on spaces
				client.Input = strings.Split(scanner.Text(), " ")

				// Determine the clients state based on the command unless
				// we're waiting for a value
				if client.State != STATE_EXPECTING_VALUE {
					// Get the command
					client.Command = client.Input[0]

					switch client.Command {
					case "get":
						client.State = STATE_COMMAND_GET
					case "set":
						client.State = STATE_COMMAND_SET
					case "delete":
						client.State = STATE_COMMAND_DELETE
					case "flush_all":
						client.State = STATE_COMMAND_FLUSH_ALL
					case "quit":
						client.State = STATE_COMMAND_QUIT
					default:
						fmt.Fprintln(connection, "ERROR")
						client.Reset()
						continue
					}
				}

				// Switch on the type of command
				switch client.State {
				// Are we expecting a value from a set command?
				case STATE_EXPECTING_VALUE:
					// If the value isn't set then set it
					if len(client.Record.Value) == 0 {
						client.Record.Value = scanner.Bytes()
					// Otherwise append to it
					} else {
						client.Record.Value = append(client.Record.Value, scanner.Bytes()...)
					}

					client.Record.Value = append(client.Record.Value, []byte{'\r', '\n'}...)

					// Count the length of the value minus the trailing \r\n
					valueLength := int64(len(client.Record.Value)) - 2

					// If the datastore is same or greater than the expected length
					// we are done with this op
					if valueLength >= client.Record.Length {
						// If it's the same length we can try and store it
						if valueLength == client.Record.Length {
							// Store the value
							datastore[client.Record.Key] = client.Record

							// Inform the client we have stored the value
							// TODO: error handling here
							fmt.Fprintln(connection, "STORED")
						// Otherwise the client has messed up
						} else {
							// Inform the client that they messed up
							fmt.Fprintln(connection, "CLIENT_ERROR")
							fmt.Fprintln(connection, valueLength)
							fmt.Fprintln(connection, "ERROR")
						}

						// Reset the clients state
						client.Reset()
					}
				// get [key1 ... keyn]
				// TODO: handling multiple key gets
				case STATE_COMMAND_GET:
					// Check if a key was passed, if so try and retrieve it
					if len(client.Input) == 2 {
						// Get the key
						key := client.Input[1]

						// Look up the record in our datastore
						record := datastore[key]

						// Did it exist?
						if record != nil {
							fmt.Fprintln(connection, fmt.Sprintf("VALUE %s %d %d", record.Key, record.Flags, record.Length))
							fmt.Fprint(connection, string(record.Value[:]))
						}

						fmt.Fprintln(connection, "END")
					} else {
						fmt.Fprintln(connection, "CLIENT_ERROR")
					}

					client.Reset()
				// set [key] [flags] [exptime] [length] [casunique] [noreply]
				case STATE_COMMAND_SET:
					// Check the right number of arguments are passed
					// casunique and noreply are optional
					if len(client.Input) == 5 {
						// Get the key name
						client.Record.Key = client.Input[1]

						// Get any flags
						client.Record.Flags, err = strconv.ParseInt(client.Input[2], 10, 64)

						if err != nil {
							fmt.Fprintln(connection, "CLIENT_ERROR ", err)
							client.Reset()
							break
						}

						// Get the key TTL
						client.Record.Ttl, err = strconv.ParseInt(client.Input[3], 10, 64)

						if err != nil {
							fmt.Fprintln(connection, "CLIENT_ERROR ", err)
							client.Reset()
							break
						}

						// Get the value length
						client.Record.Length, err = strconv.ParseInt(client.Input[4], 10, 64)

						if err != nil {
							fmt.Fprintln(connection, "CLIENT_ERROR ", err)
							client.Reset()
							break
						}

						// Set that we are expecting a value
						client.State = STATE_EXPECTING_VALUE
					} else {
						fmt.Fprintln(connection, "ERROR ", err)
						client.Reset()
					}
				// delete [key] [noreply]
				case STATE_COMMAND_DELETE:
					// Check if a key was passed, if so try and retrieve it
					if len(client.Input) == 2 {
						// Get the key
						key := client.Input[1]

						// Look up the record in our datastore
						record := datastore[key]

						// Did it exist? If so 'delete' it
						if record != nil {
							delete(datastore, key)
							fmt.Fprintln(connection, "DELETED")
						} else {
							fmt.Fprintln(connection, "NOT_FOUND")
						}

					} else {
						fmt.Fprintln(connection, "CLIENT_ERROR")
					}

					client.Reset()
				// quit
				case STATE_COMMAND_QUIT:
					// Not much to do here atm..
					// Eventually we will do logging etc
				// flush_all [delay]
				case STATE_COMMAND_FLUSH_ALL:
					// Check if a delay was passed
					if len(client.Input) == 2 {
						delay, err := strconv.ParseInt(client.Input[1], 10, 64)

						if err != nil {
							fmt.Fprintln(connection, "CLIENT_ERROR ", err)
							client.Reset()
							break
						}

						time.Sleep(time.Duration(delay) * time.Second)
					}

					// Reset the datastore
					datastore = make(map[string]*Record)
					fmt.Fprintln(connection, "OK")

					client.Reset()
				}
			}

			// Print out errors to stderr
			if err := scanner.Err(); err != nil {
				fmt.Fprintln(connection, "ERROR ", err)
				client.Reset()
			}
		}(connection)
	}
}

// Reset a clients state to what it would be on first connection
func (client *Client) Reset() {
	client.State = STATE_DEFAULT
	client.Input = []string{}
	client.Command = ""
	client.Record =  &Record{}
}
<commit_msg>The PHP extension at least seems to pass an empty space when no delay is set<commit_after>// minicache is a work in progress in-memory caching system
// featuring a similar text based protocol to memcached
package main

// TODO:
// - TTL handling
// - casunique?
// - noreply?
// - Investigate how memcached handles concurrency/locking

import (
	"bufio"
	"fmt"
	"log"
	"net"
	// "os"
	"strconv"
	"strings"
	"time"
)

type Client struct {
	Id      string
	State   uint8
	Input   []string
	Command string
	Record  *Record
}

type Record struct {
	Key    string
	Value  []byte
	Flags  int64
	Ttl    int64
	Length int64
}

const STATE_DEFAULT uint8 = 1
const STATE_EXPECTING_VALUE uint8 = 2

const STATE_COMMAND_GET uint8 = 3
const STATE_COMMAND_SET uint8 = 4
const STATE_COMMAND_DELETE uint8 = 5
const STATE_COMMAND_QUIT uint8 = 6
const STATE_COMMAND_FLUSH_ALL uint8 = 7

var ticker = time.NewTicker(time.Second * 1)
var clients map[string]*Client
var datastore map[string]*Record

func main() {
	// Start the server on port 5268
	server, err := net.Listen("tcp", ":5268")

	if err != nil {
		log.Fatal(err)
	}

	// Ensure that the server closes
	defer server.Close()

	// Our datastore..
	datastore = make(map[string]*Record)

	// A list of active clients
	clients = make(map[string]*Client)

	// Print out the datastore contents every second for debug
	go func() {
		for range ticker.C {
			fmt.Println(datastore)
		}
	}()

	for {
		// Wait for a connection.
		connection, err := server.Accept()

		if err != nil {
			log.Fatal(err)
		}

		// Handle the connection in a new goroutine.
		// The loop then returns to accepting new connections,
		// so that multiple connections may be served concurrently.
		go func(connection net.Conn) {
			// Create the client and store it
			client := &Client{
				Id:      connection.RemoteAddr().String(),
				State:   STATE_DEFAULT,
				Input:   []string{},
				Command: "",
				Record:  &Record{},
			}
			clients[connection.RemoteAddr().String()] = client

			// Ensure the client is tidied up once they're done
			defer func(connection net.Conn, clients map[string]*Client, id string) {
				delete(clients, id)
				connection.Close()
			}(connection, clients, client.Id)

			// Create a new scanner for the client input
			scanner := bufio.NewScanner(connection)

			// Handle each line (command)
			for scanner.Scan() {
				// Split the client input up based on spaces
				client.Input = strings.Split(scanner.Text(), " ")

				// Determine the clients state based on the command unless
				// we're waiting for a value
				if client.State != STATE_EXPECTING_VALUE {
					// Get the command
					client.Command = client.Input[0]

					switch client.Command {
					case "get":
						client.State = STATE_COMMAND_GET
					case "set":
						client.State = STATE_COMMAND_SET
					case "delete":
						client.State = STATE_COMMAND_DELETE
					case "flush_all":
						client.State = STATE_COMMAND_FLUSH_ALL
					case "quit":
						client.State = STATE_COMMAND_QUIT
					default:
						fmt.Fprintln(connection, "ERROR")
						client.Reset()
						continue
					}
				}

				// Switch on the type of command
				switch client.State {
				// Are we expecting a value from a set command?
				case STATE_EXPECTING_VALUE:
					// If the value isn't set then set it
					if len(client.Record.Value) == 0 {
						client.Record.Value = scanner.Bytes()
					// Otherwise append to it
					} else {
						client.Record.Value = append(client.Record.Value, scanner.Bytes()...)
					}

					client.Record.Value = append(client.Record.Value, []byte{'\r', '\n'}...)

					// Count the length of the value minus the trailing \r\n
					valueLength := int64(len(client.Record.Value)) - 2

					// If the datastore is same or greater than the expected length
					// we are done with this op
					if valueLength >= client.Record.Length {
						// If it's the same length we can try and store it
						if valueLength == client.Record.Length {
							// Store the value
							datastore[client.Record.Key] = client.Record

							// Inform the client we have stored the value
							// TODO: error handling here
							fmt.Fprintln(connection, "STORED")
						// Otherwise the client has messed up
						} else {
							// Inform the client that they messed up
							fmt.Fprintln(connection, "CLIENT_ERROR")
							fmt.Fprintln(connection, valueLength)
							fmt.Fprintln(connection, "ERROR")
						}

						// Reset the clients state
						client.Reset()
					}
				// get [key1 ... keyn]
				// TODO: handling multiple key gets
				case STATE_COMMAND_GET:
					// Check if a key was passed, if so try and retrieve it
					if len(client.Input) == 2 {
						// Get the key
						key := client.Input[1]

						// Look up the record in our datastore
						record := datastore[key]

						// Did it exist?
						if record != nil {
							fmt.Fprintln(connection, fmt.Sprintf("VALUE %s %d %d", record.Key, record.Flags, record.Length))
							fmt.Fprint(connection, string(record.Value[:]))
						}

						fmt.Fprintln(connection, "END")
					} else {
						fmt.Fprintln(connection, "CLIENT_ERROR")
					}

					client.Reset()
				// set [key] [flags] [exptime] [length] [casunique] [noreply]
				case STATE_COMMAND_SET:
					// Check the right number of arguments are passed
					// casunique and noreply are optional
					if len(client.Input) == 5 {
						// Get the key name
						client.Record.Key = client.Input[1]

						// Get any flags
						client.Record.Flags, err = strconv.ParseInt(client.Input[2], 10, 64)

						if err != nil {
							fmt.Fprintln(connection, "CLIENT_ERROR ", err)
							client.Reset()
							break
						}

						// Get the key TTL
						client.Record.Ttl, err = strconv.ParseInt(client.Input[3], 10, 64)

						if err != nil {
							fmt.Fprintln(connection, "CLIENT_ERROR ", err)
							client.Reset()
							break
						}

						// Get the value length
						client.Record.Length, err = strconv.ParseInt(client.Input[4], 10, 64)

						if err != nil {
							fmt.Fprintln(connection, "CLIENT_ERROR ", err)
							client.Reset()
							break
						}

						// Set that we are expecting a value
						client.State = STATE_EXPECTING_VALUE
					} else {
						fmt.Fprintln(connection, "ERROR ", err)
						client.Reset()
					}
				// delete [key] [noreply]
				case STATE_COMMAND_DELETE:
					// Check if a key was passed, if so try and retrieve it
					if len(client.Input) == 2 {
						// Get the key
						key := client.Input[1]

						// Look up the record in our datastore
						record := datastore[key]

						// Did it exist? If so 'delete' it
						if record != nil {
							delete(datastore, key)
							fmt.Fprintln(connection, "DELETED")
						} else {
							fmt.Fprintln(connection, "NOT_FOUND")
						}

					} else {
						fmt.Fprintln(connection, "CLIENT_ERROR")
					}

					client.Reset()
				// quit
				case STATE_COMMAND_QUIT:
					// Not much to do here atm..
					// Eventually we will do logging etc
				// flush_all [delay]
				case STATE_COMMAND_FLUSH_ALL:
					// Check if a delay was passed
					if len(client.Input) == 2 && client.Input[1] != "" {
						delay, err := strconv.ParseInt(client.Input[1], 10, 64)

						if err != nil {
							fmt.Fprintln(connection, "CLIENT_ERROR ", err)
							client.Reset()
							break
						}

						time.Sleep(time.Duration(delay) * time.Second)
					}

					// Reset the datastore
					datastore = make(map[string]*Record)
					fmt.Fprintln(connection, "OK")

					client.Reset()
				}
			}

			// Print out errors to stderr
			if err := scanner.Err(); err != nil {
				fmt.Fprintln(connection, "ERROR ", err)
				client.Reset()
			}
		}(connection)
	}
}

// Reset a clients state to what it would be on first connection
func (client *Client) Reset() {
	client.State = STATE_DEFAULT
	client.Input = []string{}
	client.Command = ""
	client.Record =  &Record{}
}
<|endoftext|>


<commit_before>// minicache is a work in progress in-memory caching system
// featuring a similar text based protocol to memcached
package main

// TODO:
// - TTL handling
// - casunique?
// - noreply?
// - Investigate how memcached handles concurrency/locking

import (
	"bufio"
	"fmt"
	"log"
	"net"
	// "os"
	"strconv"
	"strings"
	"time"
)

type Client struct {
	Id      string
	State   uint8
	Input   []string
	Command string
	Record  *Record
}

type Record struct {
	Key    string
	Value  string
	Flags  int64
	Ttl    int64
	Length int64
}

const STATE_DEFAULT uint8 = 1
const STATE_EXPECTING_VALUE uint8 = 2

const STATE_COMMAND_GET uint8 = 3
const STATE_COMMAND_SET uint8 = 4
const STATE_COMMAND_DELETE uint8 = 5
const STATE_COMMAND_QUIT uint8 = 6
const STATE_COMMAND_FLUSHALL uint8 = 7

var ticker = time.NewTicker(time.Second * 1)
var clients map[string]*Client
var datastore map[string]*Record

func main() {
	// Start the server on port 5268
	server, err := net.Listen("tcp", ":5268")

	if err != nil {
		log.Fatal(err)
	}

	// Ensure that the server closes
	defer server.Close()

	// Our datastore..
	datastore = make(map[string]*Record)

	// A list of active clients
	clients = make(map[string]*Client)

	// Print out the datastore contents every second for debug
	go func() {
		for range ticker.C {
			fmt.Println(datastore)
		}
	}()

	for {
		// Wait for a connection.
		connection, err := server.Accept()

		if err != nil {
			log.Fatal(err)
		}

		// Handle the connection in a new goroutine.
		// The loop then returns to accepting new connections,
		// so that multiple connections may be served concurrently.
		go func(connection net.Conn) {
			// Create the client and store it
			client := &Client{
				Id:      connection.RemoteAddr().String(),
				State:   STATE_DEFAULT,
				Input:   []string{},
				Command: "",
				Record:  &Record{},
			}
			clients[connection.RemoteAddr().String()] = client

			// Ensure the client is tidied up once they're done
			defer func(connection net.Conn, clients map[string]*Client, id string) {
				delete(clients, id)
				connection.Close()
			}(connection, clients, client.Id)

			// Create a new scanner for the client input
			scanner := bufio.NewScanner(connection)

			// Handle each line (command)
			for scanner.Scan() {
				// Split the client input up based on spaces
				client.Input = strings.Split(scanner.Text(), " ")

				// If we're in our default state then determine
				// what command we're running
				if client.State == STATE_DEFAULT {
					// Get the command
					client.Command = client.Input[0]

					switch client.Command {
					case "get":
						client.State = STATE_COMMAND_GET
					case "set":
						client.State = STATE_COMMAND_SET
					case "delete":
						client.State = STATE_COMMAND_DELETE
					case "flush_all":
						client.State = STATE_COMMAND_FLUSHALL
					}
				}

				// Switch on the type of command
				switch client.State {
				// Are we expecting a value from a set command?
				case STATE_EXPECTING_VALUE:
					// If the value isn't set then set it
					if client.Record.Value == "" {
						client.Record.Value = scanner.Text()
					// Otherwise append to it
					} else {
						client.Record.Value += scanner.Text()
					}

					client.Record.Value += "\r\n"

					// Count the length of the value minus the trailing \r\n
					valueLength := int64(len(client.Record.Value)) - 2

					// If the datastore is same or greater than the expected length
					// we are done with this op
					if valueLength >= client.Record.Length {
						// If it's the same length we can try and store it
						if valueLength == client.Record.Length {
							// Store the value
							datastore[client.Record.Key] = client.Record

							// Inform the client we have stored the value
							// TODO: error handling here
							fmt.Fprintln(connection, "STORED")
						// Otherwise the client has messed up
						} else {
							// Inform the client that they messed up
							fmt.Fprintln(connection, "CLIENT_ERROR")
							fmt.Fprintln(connection, valueLength)
							fmt.Fprintln(connection, "ERROR")
						}

						// Reset the clients state
						client.Reset()
					}
				// get [key1 ... keyn]
				// TODO: handling multiple key gets
				case STATE_COMMAND_GET:
					// Check if a key was passed, if so try and retrieve it
					if len(client.Input) == 2 {
						// Get the key
						key := client.Input[1]

						// Look up the record in our datastore
						record := datastore[key]

						// Did it exist?
						if record != nil {
							fmt.Fprintln(connection, fmt.Sprintf("VALUE %s %d %d", record.Key, record.Flags, record.Length))
							fmt.Fprint(connection, record.Value)
						}

						fmt.Fprintln(connection, "END")
					} else {
						fmt.Fprintln(connection, "CLIENT_ERROR")
					}

					// Reset the clients state regardless off success/failure
					client.Reset()
				// set [key] [flags] [exptime] [length] [casunique] [noreply]
				case STATE_COMMAND_SET:
					// Check the right number of arguments are passed
					if len(client.Input) == 5 {
						// Get the key name
						client.Record.Key = client.Input[1]

						// Get any flags
						client.Record.Flags, err = strconv.ParseInt(client.Input[2], 10, 64)

						if err != nil {
							fmt.Fprintln(connection, "CLIENT_ERROR ", err)
							break
						}

						// Get the key TTL
						client.Record.Ttl, err = strconv.ParseInt(client.Input[3], 10, 64)

						if err != nil {
							fmt.Fprintln(connection, "CLIENT_ERROR ", err)
							break
						}

						// Get the value length
						client.Record.Length, err = strconv.ParseInt(client.Input[4], 10, 64)

						if err != nil {
							fmt.Fprintln(connection, "CLIENT_ERROR ", err)
							break
						}

						// Set that we are expecting a value
						client.State = STATE_EXPECTING_VALUE
					} else {
						fmt.Fprintln(connection, "ERROR ", err)

						// Reset the clients state
						client.Reset()
					}
				// delete [key] [noreply]
				case STATE_COMMAND_DELETE:
					// Check if a key was passed, if so try and retrieve it
					if len(client.Input) == 2 {
						// Get the key
						key := client.Input[1]

						// Look up the record in our datastore
						record := datastore[key]

						// Did it exist? If so 'delete' it
						if record != nil {
							delete(datastore, key)
							fmt.Fprintln(connection, "DELETED")
						} else {
							fmt.Fprintln(connection, "NOT_FOUND")
						}

					} else {
						fmt.Fprintln(connection, "CLIENT_ERROR")
					}

					// Reset the clients state regardless off success/failure
					client.Reset()
				// quit
				case STATE_COMMAND_QUIT:
					// Not much to do here atm..
					// Eventually we will do logging etc
				// flushall [delay]
				// TODO: implement delay
				case STATE_COMMAND_FLUSHALL:
					datastore = make(map[string]*Record)
					fmt.Fprintln(connection, "OK")
				}
			}

			// Print out errors to stderr
			if err := scanner.Err(); err != nil {
				fmt.Fprintln(connection, "ERROR ", err)
			}
		}(connection)
	}
}

// Reset a clients state to what it would be on first connection
func (client *Client) Reset() {
	client.State = STATE_DEFAULT
	client.Input = []string{}
	client.Command = ""
	client.Record =  &Record{}
}
<commit_msg>Implement flush_all delay<commit_after>// minicache is a work in progress in-memory caching system
// featuring a similar text based protocol to memcached
package main

// TODO:
// - TTL handling
// - casunique?
// - noreply?
// - Investigate how memcached handles concurrency/locking

import (
	"bufio"
	"fmt"
	"log"
	"net"
	// "os"
	"strconv"
	"strings"
	"time"
)

type Client struct {
	Id      string
	State   uint8
	Input   []string
	Command string
	Record  *Record
}

type Record struct {
	Key    string
	Value  string
	Flags  int64
	Ttl    int64
	Length int64
}

const STATE_DEFAULT uint8 = 1
const STATE_EXPECTING_VALUE uint8 = 2

const STATE_COMMAND_GET uint8 = 3
const STATE_COMMAND_SET uint8 = 4
const STATE_COMMAND_DELETE uint8 = 5
const STATE_COMMAND_QUIT uint8 = 6
const STATE_COMMAND_FLUSHALL uint8 = 7

var ticker = time.NewTicker(time.Second * 1)
var clients map[string]*Client
var datastore map[string]*Record

func main() {
	// Start the server on port 5268
	server, err := net.Listen("tcp", ":5268")

	if err != nil {
		log.Fatal(err)
	}

	// Ensure that the server closes
	defer server.Close()

	// Our datastore..
	datastore = make(map[string]*Record)

	// A list of active clients
	clients = make(map[string]*Client)

	// Print out the datastore contents every second for debug
	go func() {
		for range ticker.C {
			fmt.Println(datastore)
		}
	}()

	for {
		// Wait for a connection.
		connection, err := server.Accept()

		if err != nil {
			log.Fatal(err)
		}

		// Handle the connection in a new goroutine.
		// The loop then returns to accepting new connections,
		// so that multiple connections may be served concurrently.
		go func(connection net.Conn) {
			// Create the client and store it
			client := &Client{
				Id:      connection.RemoteAddr().String(),
				State:   STATE_DEFAULT,
				Input:   []string{},
				Command: "",
				Record:  &Record{},
			}
			clients[connection.RemoteAddr().String()] = client

			// Ensure the client is tidied up once they're done
			defer func(connection net.Conn, clients map[string]*Client, id string) {
				delete(clients, id)
				connection.Close()
			}(connection, clients, client.Id)

			// Create a new scanner for the client input
			scanner := bufio.NewScanner(connection)

			// Handle each line (command)
			for scanner.Scan() {
				// Split the client input up based on spaces
				client.Input = strings.Split(scanner.Text(), " ")

				// If we're in our default state then determine
				// what command we're running
				if client.State == STATE_DEFAULT {
					// Get the command
					client.Command = client.Input[0]

					switch client.Command {
					case "get":
						client.State = STATE_COMMAND_GET
					case "set":
						client.State = STATE_COMMAND_SET
					case "delete":
						client.State = STATE_COMMAND_DELETE
					case "flush_all":
						client.State = STATE_COMMAND_FLUSHALL
					}
				}

				// Switch on the type of command
				switch client.State {
				// Are we expecting a value from a set command?
				case STATE_EXPECTING_VALUE:
					// If the value isn't set then set it
					if client.Record.Value == "" {
						client.Record.Value = scanner.Text()
					// Otherwise append to it
					} else {
						client.Record.Value += scanner.Text()
					}

					client.Record.Value += "\r\n"

					// Count the length of the value minus the trailing \r\n
					valueLength := int64(len(client.Record.Value)) - 2

					// If the datastore is same or greater than the expected length
					// we are done with this op
					if valueLength >= client.Record.Length {
						// If it's the same length we can try and store it
						if valueLength == client.Record.Length {
							// Store the value
							datastore[client.Record.Key] = client.Record

							// Inform the client we have stored the value
							// TODO: error handling here
							fmt.Fprintln(connection, "STORED")
						// Otherwise the client has messed up
						} else {
							// Inform the client that they messed up
							fmt.Fprintln(connection, "CLIENT_ERROR")
							fmt.Fprintln(connection, valueLength)
							fmt.Fprintln(connection, "ERROR")
						}

						// Reset the clients state
						client.Reset()
					}
				// get [key1 ... keyn]
				// TODO: handling multiple key gets
				case STATE_COMMAND_GET:
					// Check if a key was passed, if so try and retrieve it
					if len(client.Input) == 2 {
						// Get the key
						key := client.Input[1]

						// Look up the record in our datastore
						record := datastore[key]

						// Did it exist?
						if record != nil {
							fmt.Fprintln(connection, fmt.Sprintf("VALUE %s %d %d", record.Key, record.Flags, record.Length))
							fmt.Fprint(connection, record.Value)
						}

						fmt.Fprintln(connection, "END")
					} else {
						fmt.Fprintln(connection, "CLIENT_ERROR")
					}

					// Reset the clients state regardless off success/failure
					client.Reset()
				// set [key] [flags] [exptime] [length] [casunique] [noreply]
				case STATE_COMMAND_SET:
					// Check the right number of arguments are passed
					if len(client.Input) == 5 {
						// Get the key name
						client.Record.Key = client.Input[1]

						// Get any flags
						client.Record.Flags, err = strconv.ParseInt(client.Input[2], 10, 64)

						if err != nil {
							fmt.Fprintln(connection, "CLIENT_ERROR ", err)
							break
						}

						// Get the key TTL
						client.Record.Ttl, err = strconv.ParseInt(client.Input[3], 10, 64)

						if err != nil {
							fmt.Fprintln(connection, "CLIENT_ERROR ", err)
							break
						}

						// Get the value length
						client.Record.Length, err = strconv.ParseInt(client.Input[4], 10, 64)

						if err != nil {
							fmt.Fprintln(connection, "CLIENT_ERROR ", err)
							break
						}

						// Set that we are expecting a value
						client.State = STATE_EXPECTING_VALUE
					} else {
						fmt.Fprintln(connection, "ERROR ", err)

						// Reset the clients state
						client.Reset()
					}
				// delete [key] [noreply]
				case STATE_COMMAND_DELETE:
					// Check if a key was passed, if so try and retrieve it
					if len(client.Input) == 2 {
						// Get the key
						key := client.Input[1]

						// Look up the record in our datastore
						record := datastore[key]

						// Did it exist? If so 'delete' it
						if record != nil {
							delete(datastore, key)
							fmt.Fprintln(connection, "DELETED")
						} else {
							fmt.Fprintln(connection, "NOT_FOUND")
						}

					} else {
						fmt.Fprintln(connection, "CLIENT_ERROR")
					}

					// Reset the clients state regardless off success/failure
					client.Reset()
				// quit
				case STATE_COMMAND_QUIT:
					// Not much to do here atm..
					// Eventually we will do logging etc
				// flushall [delay]
				case STATE_COMMAND_FLUSHALL:
					// Check if a delay was passed
					if len(client.Input) == 2 {
						delay, err := strconv.ParseInt(client.Input[1], 10, 64)

						if err != nil {
							fmt.Fprintln(connection, "CLIENT_ERROR ", err)
							break
						}

						time.Sleep(time.Duration(delay) * time.Second)
					}

					// Reset the datastore
					datastore = make(map[string]*Record)
					fmt.Fprintln(connection, "OK")
				}
			}

			// Print out errors to stderr
			if err := scanner.Err(); err != nil {
				fmt.Fprintln(connection, "ERROR ", err)
			}
		}(connection)
	}
}

// Reset a clients state to what it would be on first connection
func (client *Client) Reset() {
	client.State = STATE_DEFAULT
	client.Input = []string{}
	client.Command = ""
	client.Record =  &Record{}
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package engine

import (
	"io"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/saltpack"
	"golang.org/x/net/context"
)

type SaltpackDecryptArg struct {
	Source io.Reader
	Sink   io.WriteCloser
	Opts   keybase1.SaltpackDecryptOptions
}

// SaltpackDecrypt decrypts data read from a source into a sink.
type SaltpackDecrypt struct {
	libkb.Contextified
	arg *SaltpackDecryptArg
	res keybase1.SaltpackEncryptedMessageInfo
}

// NewSaltpackDecrypt creates a SaltpackDecrypt engine.
func NewSaltpackDecrypt(arg *SaltpackDecryptArg, g *libkb.GlobalContext) *SaltpackDecrypt {
	return &SaltpackDecrypt{
		Contextified: libkb.NewContextified(g),
		arg:          arg,
	}
}

// Name is the unique engine name.
func (e *SaltpackDecrypt) Name() string {
	return "SaltpackDecrypt"
}

// GetPrereqs returns the engine prereqs.
func (e *SaltpackDecrypt) Prereqs() Prereqs {
	return Prereqs{}
}

// RequiredUIs returns the required UIs.
func (e *SaltpackDecrypt) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{
		libkb.SaltpackUIKind,
		libkb.SecretUIKind,
	}
}

// SubConsumers returns the other UI consumers for this engine.
func (e *SaltpackDecrypt) SubConsumers() []libkb.UIConsumer {
	return []libkb.UIConsumer{
		&SaltpackSenderIdentify{},
	}
}

func (e *SaltpackDecrypt) promptForDecrypt(ctx *Context, mki *saltpack.MessageKeyInfo) (err error) {
	defer e.G().Trace("SaltpackDecrypt::promptForDecrypt", func() error { return err })()

	spsiArg := SaltpackSenderIdentifyArg{
		isAnon:           mki.SenderIsAnon,
		publicKey:        libkb.BoxPublicKeyToKeybaseKID(mki.SenderKey),
		interactive:      e.arg.Opts.Interactive,
		forceRemoteCheck: e.arg.Opts.ForceRemoteCheck,
		reason: keybase1.IdentifyReason{
			Reason: "Identify who encrypted this message",
			Type:   keybase1.IdentifyReasonType_DECRYPT,
		},
	}

	spsiEng := NewSaltpackSenderIdentify(e.G(), &spsiArg)
	if err = RunEngine(spsiEng, ctx); err != nil {
		return err
	}

	arg := keybase1.SaltpackPromptForDecryptArg{
		Sender: spsiEng.Result(),
	}
	e.res.Sender = arg.Sender

	usedDelegateUI := false
	if e.G().UIRouter != nil {
		if ui, err := e.G().UIRouter.GetIdentifyUI(); err == nil && ui != nil {
			usedDelegateUI = true
		}
	}

	err = ctx.SaltpackUI.SaltpackPromptForDecrypt(context.TODO(), arg, usedDelegateUI)
	if err != nil {
		return err
	}
	return err
}

func (e *SaltpackDecrypt) makeMessageInfo(me *libkb.User, mki *saltpack.MessageKeyInfo) {
	if mki == nil || me == nil {
		return
	}
	ckf := me.GetComputedKeyFamily()
	for _, nr := range mki.NamedReceivers {
		kid := keybase1.KIDFromRawKey(nr, libkb.KIDNaclDH)
		if dev, _ := ckf.GetDeviceForKID(kid); dev != nil {
			edev := dev.ProtExport()
			edev.EncryptKey = kid
			e.res.Devices = append(e.res.Devices, *edev)
		}
	}
	e.res.NumAnonReceivers = mki.NumAnonReceivers
	e.res.ReceiverIsAnon = mki.ReceiverIsAnon
}

// Run starts the engine.
func (e *SaltpackDecrypt) Run(ctx *Context) (err error) {
	defer e.G().Trace("SaltpackDecrypt::Run", func() error { return err })()

	// We don't load this in the --paperkey case.
	var me *libkb.User

	var key libkb.GenericKey
	if e.arg.Opts.UsePaperKey {
		// Prompt the user for a paper key. This doesn't require you to be
		// logged in.
		keypair, _, err := getPaperKey(e.G(), ctx, nil)
		if err != nil {
			return err
		}
		key = keypair.encKey
	} else {
		// Load self so that we can get device keys. This does require you to
		// be logged in.
		me, err = libkb.LoadMe(libkb.NewLoadUserArg(e.G()))
		if err != nil {
			return err
		}
		// Get the device encryption key, maybe prompting the user.
		ska := libkb.SecretKeyArg{
			Me:      me,
			KeyType: libkb.DeviceEncryptionKeyType,
		}
		e.G().Log.Debug("| GetSecretKeyWithPrompt")
		key, err = e.G().Keyrings.GetSecretKeyWithPrompt(ctx.SecretKeyPromptArg(ska, "decrypting a message/file"))
		if err != nil {
			return err
		}
	}

	kp, ok := key.(libkb.NaclDHKeyPair)
	if !ok || kp.Private == nil {
		return libkb.KeyCannotDecryptError{}
	}

	// For DH mode.
	hookMki := func(mki *saltpack.MessageKeyInfo) error {
		return e.promptForDecrypt(ctx, mki)
	}

	// For signcryption mode.
	hookSenderSigningKey := func(senderSigningKey saltpack.SigningPublicKey) error {
		// TODO: implement sender checking/prompting
		return nil
	}

	e.G().Log.Debug("| SaltpackDecrypt")
	var mki *saltpack.MessageKeyInfo
	mki, err = libkb.SaltpackDecrypt(e.G(), e.arg.Source, e.arg.Sink, kp, hookMki, hookSenderSigningKey)
	if err == saltpack.ErrNoDecryptionKey {
		err = libkb.NoDecryptionKeyError{Msg: "no suitable device key found"}
	}

	// It's ok if me is nil here.
	e.makeMessageInfo(me, mki)

	return err
}

func (e *SaltpackDecrypt) MessageInfo() keybase1.SaltpackEncryptedMessageInfo {
	return e.res
}
<commit_msg>plumb promptForDecrypt for the signcryption mode<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package engine

import (
	"bytes"
	"io"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/saltpack"
	"golang.org/x/net/context"
)

type SaltpackDecryptArg struct {
	Source io.Reader
	Sink   io.WriteCloser
	Opts   keybase1.SaltpackDecryptOptions
}

// SaltpackDecrypt decrypts data read from a source into a sink.
type SaltpackDecrypt struct {
	libkb.Contextified
	arg *SaltpackDecryptArg
	res keybase1.SaltpackEncryptedMessageInfo
}

// NewSaltpackDecrypt creates a SaltpackDecrypt engine.
func NewSaltpackDecrypt(arg *SaltpackDecryptArg, g *libkb.GlobalContext) *SaltpackDecrypt {
	return &SaltpackDecrypt{
		Contextified: libkb.NewContextified(g),
		arg:          arg,
	}
}

// Name is the unique engine name.
func (e *SaltpackDecrypt) Name() string {
	return "SaltpackDecrypt"
}

// GetPrereqs returns the engine prereqs.
func (e *SaltpackDecrypt) Prereqs() Prereqs {
	return Prereqs{}
}

// RequiredUIs returns the required UIs.
func (e *SaltpackDecrypt) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{
		libkb.SaltpackUIKind,
		libkb.SecretUIKind,
	}
}

// SubConsumers returns the other UI consumers for this engine.
func (e *SaltpackDecrypt) SubConsumers() []libkb.UIConsumer {
	return []libkb.UIConsumer{
		&SaltpackSenderIdentify{},
	}
}

func (e *SaltpackDecrypt) promptForDecrypt(ctx *Context, publicKey keybase1.KID, isAnon bool) (err error) {
	defer e.G().Trace("SaltpackDecrypt::promptForDecrypt", func() error { return err })()

	spsiArg := SaltpackSenderIdentifyArg{
		isAnon:           isAnon,
		publicKey:        publicKey,
		interactive:      e.arg.Opts.Interactive,
		forceRemoteCheck: e.arg.Opts.ForceRemoteCheck,
		reason: keybase1.IdentifyReason{
			Reason: "Identify who encrypted this message",
			Type:   keybase1.IdentifyReasonType_DECRYPT,
		},
	}

	spsiEng := NewSaltpackSenderIdentify(e.G(), &spsiArg)
	if err = RunEngine(spsiEng, ctx); err != nil {
		return err
	}

	arg := keybase1.SaltpackPromptForDecryptArg{
		Sender: spsiEng.Result(),
	}
	e.res.Sender = arg.Sender

	usedDelegateUI := false
	if e.G().UIRouter != nil {
		if ui, err := e.G().UIRouter.GetIdentifyUI(); err == nil && ui != nil {
			usedDelegateUI = true
		}
	}

	err = ctx.SaltpackUI.SaltpackPromptForDecrypt(context.TODO(), arg, usedDelegateUI)
	if err != nil {
		return err
	}
	return err
}

func (e *SaltpackDecrypt) makeMessageInfo(me *libkb.User, mki *saltpack.MessageKeyInfo) {
	if mki == nil || me == nil {
		return
	}
	ckf := me.GetComputedKeyFamily()
	for _, nr := range mki.NamedReceivers {
		kid := keybase1.KIDFromRawKey(nr, libkb.KIDNaclDH)
		if dev, _ := ckf.GetDeviceForKID(kid); dev != nil {
			edev := dev.ProtExport()
			edev.EncryptKey = kid
			e.res.Devices = append(e.res.Devices, *edev)
		}
	}
	e.res.NumAnonReceivers = mki.NumAnonReceivers
	e.res.ReceiverIsAnon = mki.ReceiverIsAnon
}

// Run starts the engine.
func (e *SaltpackDecrypt) Run(ctx *Context) (err error) {
	defer e.G().Trace("SaltpackDecrypt::Run", func() error { return err })()

	// We don't load this in the --paperkey case.
	var me *libkb.User

	var key libkb.GenericKey
	if e.arg.Opts.UsePaperKey {
		// Prompt the user for a paper key. This doesn't require you to be
		// logged in.
		keypair, _, err := getPaperKey(e.G(), ctx, nil)
		if err != nil {
			return err
		}
		key = keypair.encKey
	} else {
		// Load self so that we can get device keys. This does require you to
		// be logged in.
		me, err = libkb.LoadMe(libkb.NewLoadUserArg(e.G()))
		if err != nil {
			return err
		}
		// Get the device encryption key, maybe prompting the user.
		ska := libkb.SecretKeyArg{
			Me:      me,
			KeyType: libkb.DeviceEncryptionKeyType,
		}
		e.G().Log.Debug("| GetSecretKeyWithPrompt")
		key, err = e.G().Keyrings.GetSecretKeyWithPrompt(ctx.SecretKeyPromptArg(ska, "decrypting a message/file"))
		if err != nil {
			return err
		}
	}

	kp, ok := key.(libkb.NaclDHKeyPair)
	if !ok || kp.Private == nil {
		return libkb.KeyCannotDecryptError{}
	}

	// For DH mode.
	hookMki := func(mki *saltpack.MessageKeyInfo) error {
		kidToIdentify := libkb.BoxPublicKeyToKeybaseKID(mki.SenderKey)
		return e.promptForDecrypt(ctx, kidToIdentify, mki.SenderIsAnon)
	}

	// For signcryption mode.
	hookSenderSigningKey := func(senderSigningKey saltpack.SigningPublicKey) error {
		kidToIdentify := libkb.SigningPublicKeyToKeybaseKID(senderSigningKey)
		// See if the sender signing key is all zeroes.
		isAnon := bytes.Equal(senderSigningKey.ToKID(), make([]byte, len(senderSigningKey.ToKID())))
		return e.promptForDecrypt(ctx, kidToIdentify, isAnon)
	}

	e.G().Log.Debug("| SaltpackDecrypt")
	var mki *saltpack.MessageKeyInfo
	mki, err = libkb.SaltpackDecrypt(e.G(), e.arg.Source, e.arg.Sink, kp, hookMki, hookSenderSigningKey)
	if err == saltpack.ErrNoDecryptionKey {
		err = libkb.NoDecryptionKeyError{Msg: "no suitable device key found"}
	}

	// It's ok if me is nil here.
	e.makeMessageInfo(me, mki)

	return err
}

func (e *SaltpackDecrypt) MessageInfo() keybase1.SaltpackEncryptedMessageInfo {
	return e.res
}
<|endoftext|>


<commit_before>package pdb

import (
	"fmt"

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

type Entry struct {
	Path   string
	IdCode string
	Chains []*Chain
}

type Chain struct {
	Entry    *Entry
	Ident    byte
	SeqType  SequenceType
	Sequence []seq.Residue
	Models   []*Model
	Missing  []*Residue
}

type Model struct {
	Entry    *Entry
	Chain    *Chain
	Num      int
	Residues []*Residue
}

type Residue struct {
	Name          seq.Residue
	SequenceNum   int
	InsertionCode byte
	Atoms         []Atom
}

type Atom struct {
	Name string
	Het  bool
	Coords
}

type Coords struct {
	X, Y, Z float64
}

// Chain returns a chain with the given identifier.
// If such a chain does not exist, nil is returned.
func (entry *Entry) Chain(ident byte) *Chain {
	for i := range entry.Chains {
		if entry.Chains[i].Ident == ident {
			return entry.Chains[i]
		}
	}
	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 (entry *Entry) OneChain() *Chain {
	if len(entry.Chains) != 1 {
		panic(fmt.Sprintf("OneChain can only be called on PDB entries with "+
			"ONE chain. But the '%s' PDB entry has %d chains.",
			entry.Path, len(entry.Chains)))
	}
	return entry.Chains[0]
}

// IsProtein returns true if the chain consists of amino acids.
func (c Chain) IsProtein() bool {
	return c.SeqType == SeqProtein && len(c.Models) > 0
}

// SequenceCaAtomSlice attempts to extract a contiguous slice of alpha-carbon
// ATOM records based on *residue* index. Namely, if a contiguous slice cannot
// be found, nil is returned. If there is more than one model, the first model
// is used.
func (c Chain) SequenceCaAtomSlice(start, end int) []Coords {
	return c.Models[0].SequenceCaAtomSlice(start, end)
}

// SequenceCaAtoms returns a slice of all Ca atoms for the chain in
// correspondence with the sequence in SEQRES (automatically using the first
// model).
//
// See Model.SequenceCaAtoms for the deets.
func (c Chain) SequenceCaAtoms() ([]*Coords, error) {
	return c.Models[0].SequenceCaAtoms()
}

// SequenceAtoms returns a slice of all residues for the chain in
// correspondence with the sequence in SEQRES (automatically using the first
// model). Namely, the mapping is sparse, since not all SEQRES residues have
// an ATOM record.
//
// See Model.SequenceCaAtoms for the deets.
func (c Chain) SequenceAtoms() ([]*Residue, error) {
	return c.Models[0].SequenceAtoms()
}

// CaAtoms returns all alpha-carbon atoms in the chain. If there is more than
// one model, only the first model is used.
func (c Chain) CaAtoms() []Coords {
	return c.Models[0].CaAtoms()
}

// SequenceCaAtomSlice attempts to extract a contiguous slice of alpha-carbon
// ATOM records based on *residue* index. Namely, if a contiguous slice cannot
// be found, nil is returned.
func (m Model) SequenceCaAtomSlice(start, end int) []Coords {
	residues, err := m.SequenceCaAtoms()
	if err != nil {
		fmt.Printf("%s\n", err)
		return nil
	}

	atoms := make([]Coords, end-start)
	for i, cai := 0, start; cai < end; i, cai = i+1, cai+1 {
		if residues[start] == nil {
			return nil
		}
		atoms[i] = *residues[cai]
	}
	return atoms
}

// SequenceCaAtoms returns a slice of all Ca atoms for the model in
// correspondence with the sequence in SEQRES.
// Note that a slice of pointers is returned, since not all residues
// necessarily correspond to a alpha-carbon ATOM.
//
// This method can proceed in one of two ways. First, if "REMARK 465" is
// present in the PDB file, it will be used to determine the positions of the
// holes in the sequence (i.e., residues in SEQRES without an ATOM record).
// This method is generally reliable, since REMARK 465 lists all residues
// in SEQRES that don't have an ATOM record. This will fail if there are any
// unreported missing residues.
//
// If "REMARK 465" is absent, then we have to rely on the order of ATOM records
// to correspond to a residue index in the SEQRES sequence. This will fail
// with an error if there are any unreported missing residues.
//
// Generally, false positives are limited by returning errors if corruption
// is detected. However, false positives can be returned in pathological cases
// (like long strings of low complexity regions or UNKNOWN amino acids), but
// they are rare. Probably on the order of a handful in the entire PDB.
//
// In sum, a list of atom pointers is returned with length equal to the number
// of residues in the SEQRES record for this model. Some pointers may be nil.
func (m Model) SequenceCaAtoms() ([]*Coords, error) {
	mapping, err := m.SequenceAtoms()
	if err != nil {
		return nil, err
	}

	cas := make([]*Coords, len(mapping))
	for i := range mapping {
		if mapping[i] != nil {
			cas[i] = mapping[i].Ca()
		}
	}
	return cas, nil
}

// SequenceAtoms is just like SequenceCaAtoms, except it returns the residues
// instead of the alpha-carbon coordinates directly. The advantage here is to
// get a mapping that isn't limited by the presence of alpha-carbon atoms.
//
// See SequenceCaAtoms for the deets.
func (m Model) SequenceAtoms() ([]*Residue, error) {
	if len(m.Chain.Missing) > 0 {
		return m.seqAtomsChunksMerge()
	}
	return m.seqAtomsGuess()
}

// CaAtoms returns all alpha-carbon atoms in the model.
// This includes multiple alpha-carbon atoms belonging to the same residue.
// It does not include HETATMs.
func (m Model) CaAtoms() []Coords {
	cas := make([]Coords, 0, len(m.Residues))
	for _, r := range m.Residues {
		for _, atom := range r.Atoms {
			if atom.Name == "CA" && !atom.Het {
				cas = append(cas, atom.Coords)
			}
		}
	}
	return cas
}

// Ca returns the alpha-carbon atom in this residue.
// If one does not exist, nil is returned.
func (r Residue) Ca() *Coords {
	for _, atom := range r.Atoms {
		if atom.Name == "CA" {
			return &atom.Coords
		}
	}
	return nil
}

func (coords Coords) String() string {
	return fmt.Sprintf("%0.3f %0.3f %0.3f", coords.X, coords.Y, coords.Z)
}
<commit_msg>Don't check the same index every time.<commit_after>package pdb

import (
	"fmt"

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

type Entry struct {
	Path   string
	IdCode string
	Chains []*Chain
}

type Chain struct {
	Entry    *Entry
	Ident    byte
	SeqType  SequenceType
	Sequence []seq.Residue
	Models   []*Model
	Missing  []*Residue
}

type Model struct {
	Entry    *Entry
	Chain    *Chain
	Num      int
	Residues []*Residue
}

type Residue struct {
	Name          seq.Residue
	SequenceNum   int
	InsertionCode byte
	Atoms         []Atom
}

type Atom struct {
	Name string
	Het  bool
	Coords
}

type Coords struct {
	X, Y, Z float64
}

// Chain returns a chain with the given identifier.
// If such a chain does not exist, nil is returned.
func (entry *Entry) Chain(ident byte) *Chain {
	for i := range entry.Chains {
		if entry.Chains[i].Ident == ident {
			return entry.Chains[i]
		}
	}
	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 (entry *Entry) OneChain() *Chain {
	if len(entry.Chains) != 1 {
		panic(fmt.Sprintf("OneChain can only be called on PDB entries with "+
			"ONE chain. But the '%s' PDB entry has %d chains.",
			entry.Path, len(entry.Chains)))
	}
	return entry.Chains[0]
}

// IsProtein returns true if the chain consists of amino acids.
func (c Chain) IsProtein() bool {
	return c.SeqType == SeqProtein && len(c.Models) > 0
}

// SequenceCaAtomSlice attempts to extract a contiguous slice of alpha-carbon
// ATOM records based on *residue* index. Namely, if a contiguous slice cannot
// be found, nil is returned. If there is more than one model, the first model
// is used.
func (c Chain) SequenceCaAtomSlice(start, end int) []Coords {
	return c.Models[0].SequenceCaAtomSlice(start, end)
}

// SequenceCaAtoms returns a slice of all Ca atoms for the chain in
// correspondence with the sequence in SEQRES (automatically using the first
// model).
//
// See Model.SequenceCaAtoms for the deets.
func (c Chain) SequenceCaAtoms() ([]*Coords, error) {
	return c.Models[0].SequenceCaAtoms()
}

// SequenceAtoms returns a slice of all residues for the chain in
// correspondence with the sequence in SEQRES (automatically using the first
// model). Namely, the mapping is sparse, since not all SEQRES residues have
// an ATOM record.
//
// See Model.SequenceCaAtoms for the deets.
func (c Chain) SequenceAtoms() ([]*Residue, error) {
	return c.Models[0].SequenceAtoms()
}

// CaAtoms returns all alpha-carbon atoms in the chain. If there is more than
// one model, only the first model is used.
func (c Chain) CaAtoms() []Coords {
	return c.Models[0].CaAtoms()
}

// SequenceCaAtomSlice attempts to extract a contiguous slice of alpha-carbon
// ATOM records based on *residue* index. Namely, if a contiguous slice cannot
// be found, nil is returned.
func (m Model) SequenceCaAtomSlice(start, end int) []Coords {
	residues, err := m.SequenceCaAtoms()
	if err != nil {
		fmt.Printf("%s\n", err)
		return nil
	}

	atoms := make([]Coords, end-start)
	for i, cai := 0, start; cai < end; i, cai = i+1, cai+1 {
		if residues[cai] == nil {
			return nil
		}
		atoms[i] = *residues[cai]
	}
	return atoms
}

// SequenceCaAtoms returns a slice of all Ca atoms for the model in
// correspondence with the sequence in SEQRES.
// Note that a slice of pointers is returned, since not all residues
// necessarily correspond to a alpha-carbon ATOM.
//
// This method can proceed in one of two ways. First, if "REMARK 465" is
// present in the PDB file, it will be used to determine the positions of the
// holes in the sequence (i.e., residues in SEQRES without an ATOM record).
// This method is generally reliable, since REMARK 465 lists all residues
// in SEQRES that don't have an ATOM record. This will fail if there are any
// unreported missing residues.
//
// If "REMARK 465" is absent, then we have to rely on the order of ATOM records
// to correspond to a residue index in the SEQRES sequence. This will fail
// with an error if there are any unreported missing residues.
//
// Generally, false positives are limited by returning errors if corruption
// is detected. However, false positives can be returned in pathological cases
// (like long strings of low complexity regions or UNKNOWN amino acids), but
// they are rare. Probably on the order of a handful in the entire PDB.
//
// In sum, a list of atom pointers is returned with length equal to the number
// of residues in the SEQRES record for this model. Some pointers may be nil.
func (m Model) SequenceCaAtoms() ([]*Coords, error) {
	mapping, err := m.SequenceAtoms()
	if err != nil {
		return nil, err
	}

	cas := make([]*Coords, len(mapping))
	for i := range mapping {
		if mapping[i] != nil {
			cas[i] = mapping[i].Ca()
		}
	}
	return cas, nil
}

// SequenceAtoms is just like SequenceCaAtoms, except it returns the residues
// instead of the alpha-carbon coordinates directly. The advantage here is to
// get a mapping that isn't limited by the presence of alpha-carbon atoms.
//
// See SequenceCaAtoms for the deets.
func (m Model) SequenceAtoms() ([]*Residue, error) {
	if len(m.Chain.Missing) > 0 {
		return m.seqAtomsChunksMerge()
	}
	return m.seqAtomsGuess()
}

// CaAtoms returns all alpha-carbon atoms in the model.
// This includes multiple alpha-carbon atoms belonging to the same residue.
// It does not include HETATMs.
func (m Model) CaAtoms() []Coords {
	cas := make([]Coords, 0, len(m.Residues))
	for _, r := range m.Residues {
		for _, atom := range r.Atoms {
			if atom.Name == "CA" && !atom.Het {
				cas = append(cas, atom.Coords)
			}
		}
	}
	return cas
}

// Ca returns the alpha-carbon atom in this residue.
// If one does not exist, nil is returned.
func (r Residue) Ca() *Coords {
	for _, atom := range r.Atoms {
		if atom.Name == "CA" {
			return &atom.Coords
		}
	}
	return nil
}

func (coords Coords) String() string {
	return fmt.Sprintf("%0.3f %0.3f %0.3f", coords.X, coords.Y, coords.Z)
}
<|endoftext|>