<commit_before>package wallet

import (
	"errors"
	"os"
	"path/filepath"
	"sync"

	"github.com/btcsuite/btcd/chaincfg"
	"github.com/btcsuite/btcutil/hdkeychain"
	"github.com/btcsuite/btcwallet/internal/prompt"
	"github.com/btcsuite/btcwallet/waddrmgr"
	"github.com/btcsuite/btcwallet/walletdb"
	"github.com/btcsuite/btcwallet/wtxmgr"
)

const (
	walletDbName = "wallet.db"
)

var (
	// ErrLoaded describes the error condition of attempting to load or
	// create a wallet when the loader has already done so.
	ErrLoaded = errors.New("wallet already loaded")

	// ErrExists describes the error condition of attempting to create a new
	// wallet when one exists already.
	ErrExists = errors.New("wallet already exists")
)

// Loader implements the creating of new and opening of existing wallets, while
// providing a callback system for other subsystems to handle the loading of a
// wallet.  This is primarely intended for use by the RPC servers, to enable
// methods and services which require the wallet when the wallet is loaded by
// another subsystem.
//
// Loader is safe for concurrent access.
type Loader struct {
	callbacks   []func(*Wallet, walletdb.DB)
	chainParams *chaincfg.Params
	dbDirPath   string
	wallet      *Wallet
	db          walletdb.DB
	mu          sync.Mutex
}

// NewLoader constructs a Loader.
func NewLoader(chainParams *chaincfg.Params, dbDirPath string) *Loader {
	return &Loader{
		chainParams: chainParams,
		dbDirPath:   dbDirPath,
	}
}

// onLoaded executes each added callback and prevents loader from loading any
// additional wallets.  Requires mutex to be locked.
func (l *Loader) onLoaded(w *Wallet, db walletdb.DB) {
	for _, fn := range l.callbacks {
		fn(w, db)
	}

	l.wallet = w
	l.db = db
	l.callbacks = nil // not needed anymore
}

// RunAfterLoad adds a function to be executed when the loader creates or opens
// a wallet.  Functions are executed in a single goroutine in the order they are
// added.
func (l *Loader) RunAfterLoad(fn func(*Wallet, walletdb.DB)) {
	l.mu.Lock()
	if l.wallet != nil {
		w := l.wallet
		db := l.db
		l.mu.Unlock()
		fn(w, db)
	} else {
		l.callbacks = append(l.callbacks, fn)
		l.mu.Unlock()
	}
}

// CreateNewWallet creates a new wallet using the provided public and private
// passphrases.  The seed is optional.  If non-nil, addresses are derived from
// this seed.  If nil, a secure random seed is generated.
func (l *Loader) CreateNewWallet(pubPassphrase, privPassphrase, seed []byte) (*Wallet, error) {
	defer l.mu.Unlock()
	l.mu.Lock()

	if l.wallet != nil {
		return nil, ErrLoaded
	}

	dbPath := filepath.Join(l.dbDirPath, walletDbName)
	exists, err := fileExists(dbPath)
	if err != nil {
		return nil, err
	}
	if exists {
		return nil, ErrExists
	}

	// Create the wallet database backed by bolt db.
	err = os.MkdirAll(l.dbDirPath, 0700)
	if err != nil {
		return nil, err
	}
	db, err := walletdb.Create("bdb", dbPath)
	if err != nil {
		return nil, err
	}

	// Create the address manager.
	if seed != nil {
		if len(seed) < hdkeychain.MinSeedBytes ||
			len(seed) > hdkeychain.MaxSeedBytes {

			return nil, hdkeychain.ErrInvalidSeedLen
		}
	}
	addrMgrNamespace, err := db.Namespace(waddrmgrNamespaceKey)
	if err != nil {
		return nil, err
	}
	_, err = waddrmgr.Create(addrMgrNamespace, seed, pubPassphrase,
		privPassphrase, l.chainParams, nil)
	if err != nil {
		return nil, err
	}

	// Create empty transaction manager.
	txMgrNamespace, err := db.Namespace(wtxmgrNamespaceKey)
	if err != nil {
		return nil, err
	}
	_, err = wtxmgr.Create(txMgrNamespace)
	if err != nil {
		return nil, err
	}

	// Open the newly-created wallet.
	w, err := Open(pubPassphrase, l.chainParams, db, addrMgrNamespace, txMgrNamespace, nil)
	if err != nil {
		return nil, err
	}

	l.onLoaded(w, db)
	return w, nil
}

var errNoConsole = errors.New("db upgrade requires console access for additional input")

func noConsole() ([]byte, error) {
	return nil, errNoConsole
}

// OpenExistingWallet opens the wallet from the loader's wallet database path
// and the public passphrase.  If the loader is being called by a context where
// standard input prompts may be used during wallet upgrades, setting
// canConsolePrompt will enables these prompts.
func (l *Loader) OpenExistingWallet(pubPassphrase []byte, canConsolePrompt bool) (*Wallet, error) {
	defer l.mu.Unlock()
	l.mu.Lock()

	if l.wallet != nil {
		return nil, ErrLoaded
	}

	// Ensure that the network directory exists.
	if err := checkCreateDir(l.dbDirPath); err != nil {
		return nil, err
	}

	// Open the database using the boltdb backend.
	dbPath := filepath.Join(l.dbDirPath, walletDbName)
	db, err := walletdb.Open("bdb", dbPath)
	if err != nil {
		log.Errorf("Failed to open database: %v", err)
		return nil, err
	}

	addrMgrNS, err := db.Namespace(waddrmgrNamespaceKey)
	if err != nil {
		return nil, err
	}
	txMgrNS, err := db.Namespace(wtxmgrNamespaceKey)
	if err != nil {
		return nil, err
	}
	var cbs *waddrmgr.OpenCallbacks
	if canConsolePrompt {
		cbs = &waddrmgr.OpenCallbacks{
			ObtainSeed:        prompt.ProvideSeed,
			ObtainPrivatePass: prompt.ProvidePrivPassphrase,
		}
	} else {
		cbs = &waddrmgr.OpenCallbacks{
			ObtainSeed:        noConsole,
			ObtainPrivatePass: noConsole,
		}
	}
	w, err := Open(pubPassphrase, l.chainParams, db, addrMgrNS, txMgrNS, cbs)
	if err != nil {
		return nil, err
	}
	w.Start()

	l.onLoaded(w, db)
	return w, nil
}

// WalletExists returns whether a file exists at the loader's database path.
// This may return an error for unexpected I/O failures.
func (l *Loader) WalletExists() (bool, error) {
	dbPath := filepath.Join(l.dbDirPath, walletDbName)
	return fileExists(dbPath)
}

// LoadedWallet returns the loaded wallet, if any, and a bool for whether the
// wallet has been loaded or not.  If true, the wallet pointer should be safe to
// dereference.
func (l *Loader) LoadedWallet() (*Wallet, bool) {
	l.mu.Lock()
	w := l.wallet
	l.mu.Unlock()
	return w, w != nil
}

func fileExists(filePath string) (bool, error) {
	_, err := os.Stat(filePath)
	if err != nil {
		if os.IsNotExist(err) {
			return false, nil
		}
		return false, err
	}
	return true, nil
}
<commit_msg>Generate a seed when creating a new wallet if none was provided.<commit_after>package wallet

import (
	"errors"
	"os"
	"path/filepath"
	"sync"

	"github.com/btcsuite/btcd/chaincfg"
	"github.com/btcsuite/btcutil/hdkeychain"
	"github.com/btcsuite/btcwallet/internal/prompt"
	"github.com/btcsuite/btcwallet/waddrmgr"
	"github.com/btcsuite/btcwallet/walletdb"
	"github.com/btcsuite/btcwallet/wtxmgr"
)

const (
	walletDbName = "wallet.db"
)

var (
	// ErrLoaded describes the error condition of attempting to load or
	// create a wallet when the loader has already done so.
	ErrLoaded = errors.New("wallet already loaded")

	// ErrExists describes the error condition of attempting to create a new
	// wallet when one exists already.
	ErrExists = errors.New("wallet already exists")
)

// Loader implements the creating of new and opening of existing wallets, while
// providing a callback system for other subsystems to handle the loading of a
// wallet.  This is primarely intended for use by the RPC servers, to enable
// methods and services which require the wallet when the wallet is loaded by
// another subsystem.
//
// Loader is safe for concurrent access.
type Loader struct {
	callbacks   []func(*Wallet, walletdb.DB)
	chainParams *chaincfg.Params
	dbDirPath   string
	wallet      *Wallet
	db          walletdb.DB
	mu          sync.Mutex
}

// NewLoader constructs a Loader.
func NewLoader(chainParams *chaincfg.Params, dbDirPath string) *Loader {
	return &Loader{
		chainParams: chainParams,
		dbDirPath:   dbDirPath,
	}
}

// onLoaded executes each added callback and prevents loader from loading any
// additional wallets.  Requires mutex to be locked.
func (l *Loader) onLoaded(w *Wallet, db walletdb.DB) {
	for _, fn := range l.callbacks {
		fn(w, db)
	}

	l.wallet = w
	l.db = db
	l.callbacks = nil // not needed anymore
}

// RunAfterLoad adds a function to be executed when the loader creates or opens
// a wallet.  Functions are executed in a single goroutine in the order they are
// added.
func (l *Loader) RunAfterLoad(fn func(*Wallet, walletdb.DB)) {
	l.mu.Lock()
	if l.wallet != nil {
		w := l.wallet
		db := l.db
		l.mu.Unlock()
		fn(w, db)
	} else {
		l.callbacks = append(l.callbacks, fn)
		l.mu.Unlock()
	}
}

// CreateNewWallet creates a new wallet using the provided public and private
// passphrases.  The seed is optional.  If non-nil, addresses are derived from
// this seed.  If nil, a secure random seed is generated.
func (l *Loader) CreateNewWallet(pubPassphrase, privPassphrase, seed []byte) (*Wallet, error) {
	defer l.mu.Unlock()
	l.mu.Lock()

	if l.wallet != nil {
		return nil, ErrLoaded
	}

	dbPath := filepath.Join(l.dbDirPath, walletDbName)
	exists, err := fileExists(dbPath)
	if err != nil {
		return nil, err
	}
	if exists {
		return nil, ErrExists
	}

	// Create the wallet database backed by bolt db.
	err = os.MkdirAll(l.dbDirPath, 0700)
	if err != nil {
		return nil, err
	}
	db, err := walletdb.Create("bdb", dbPath)
	if err != nil {
		return nil, err
	}

	// If a seed was provided, ensure that it is of valid length. Otherwise,
	// we generate a random seed for the wallet with the recommended seed
	// length.
	if seed != nil {
		if len(seed) < hdkeychain.MinSeedBytes ||
			len(seed) > hdkeychain.MaxSeedBytes {

			return nil, hdkeychain.ErrInvalidSeedLen
		}
	} else {
		hdSeed, err := hdkeychain.GenerateSeed(hdkeychain.RecommendedSeedLen)
		if err != nil {
			return nil, err
		}
		seed = hdSeed
	}

	// Create the address manager.
	addrMgrNamespace, err := db.Namespace(waddrmgrNamespaceKey)
	if err != nil {
		return nil, err
	}
	_, err = waddrmgr.Create(addrMgrNamespace, seed, pubPassphrase,
		privPassphrase, l.chainParams, nil)
	if err != nil {
		return nil, err
	}

	// Create empty transaction manager.
	txMgrNamespace, err := db.Namespace(wtxmgrNamespaceKey)
	if err != nil {
		return nil, err
	}
	_, err = wtxmgr.Create(txMgrNamespace)
	if err != nil {
		return nil, err
	}

	// Open the newly-created wallet.
	w, err := Open(pubPassphrase, l.chainParams, db, addrMgrNamespace, txMgrNamespace, nil)
	if err != nil {
		return nil, err
	}

	l.onLoaded(w, db)
	return w, nil
}

var errNoConsole = errors.New("db upgrade requires console access for additional input")

func noConsole() ([]byte, error) {
	return nil, errNoConsole
}

// OpenExistingWallet opens the wallet from the loader's wallet database path
// and the public passphrase.  If the loader is being called by a context where
// standard input prompts may be used during wallet upgrades, setting
// canConsolePrompt will enables these prompts.
func (l *Loader) OpenExistingWallet(pubPassphrase []byte, canConsolePrompt bool) (*Wallet, error) {
	defer l.mu.Unlock()
	l.mu.Lock()

	if l.wallet != nil {
		return nil, ErrLoaded
	}

	// Ensure that the network directory exists.
	if err := checkCreateDir(l.dbDirPath); err != nil {
		return nil, err
	}

	// Open the database using the boltdb backend.
	dbPath := filepath.Join(l.dbDirPath, walletDbName)
	db, err := walletdb.Open("bdb", dbPath)
	if err != nil {
		log.Errorf("Failed to open database: %v", err)
		return nil, err
	}

	addrMgrNS, err := db.Namespace(waddrmgrNamespaceKey)
	if err != nil {
		return nil, err
	}
	txMgrNS, err := db.Namespace(wtxmgrNamespaceKey)
	if err != nil {
		return nil, err
	}
	var cbs *waddrmgr.OpenCallbacks
	if canConsolePrompt {
		cbs = &waddrmgr.OpenCallbacks{
			ObtainSeed:        prompt.ProvideSeed,
			ObtainPrivatePass: prompt.ProvidePrivPassphrase,
		}
	} else {
		cbs = &waddrmgr.OpenCallbacks{
			ObtainSeed:        noConsole,
			ObtainPrivatePass: noConsole,
		}
	}
	w, err := Open(pubPassphrase, l.chainParams, db, addrMgrNS, txMgrNS, cbs)
	if err != nil {
		return nil, err
	}
	w.Start()

	l.onLoaded(w, db)
	return w, nil
}

// WalletExists returns whether a file exists at the loader's database path.
// This may return an error for unexpected I/O failures.
func (l *Loader) WalletExists() (bool, error) {
	dbPath := filepath.Join(l.dbDirPath, walletDbName)
	return fileExists(dbPath)
}

// LoadedWallet returns the loaded wallet, if any, and a bool for whether the
// wallet has been loaded or not.  If true, the wallet pointer should be safe to
// dereference.
func (l *Loader) LoadedWallet() (*Wallet, bool) {
	l.mu.Lock()
	w := l.wallet
	l.mu.Unlock()
	return w, w != nil
}

func fileExists(filePath string) (bool, error) {
	_, err := os.Stat(filePath)
	if err != nil {
		if os.IsNotExist(err) {
			return false, nil
		}
		return false, err
	}
	return true, nil
}
<|endoftext|>


<commit_before>package server

import (
	"crypto/rand"
	"encoding/base64"
	goerr "errors"
	"fmt"
	"log"
	"net/http"

	"github.com/coreos/go-oidc"
	logrus "github.com/sirupsen/logrus"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
)

// various oidc invalid argument errors. Use 'goerror' instead of internal
// 'errors' library b/c stack trace isn't useful
var (
	notConfigured = goerr.New("OIDC ID provider configuration not found")
)

var tokenChan chan tokenInfo

type sessionInfo struct {
	Nonce string
	Token string
}

var stateInfoMap map[string]sessionInfo

func init() {
	tokenChan = make(chan tokenInfo)
	stateInfoMap = make(map[string]sessionInfo)
}

// InternalOIDCProvider is our internal representation of an oidc provider
type InternalOIDCProvider struct {
	Provider     *oidc.Provider
	Issuer       string
	ClientID     string
	ClientSecret string
	RedirectURI  string
}

// CryptoString returns a cryptographically random, URL safe string with length at least n
func CryptoString(n int) string {
	var numBytes int
	for n >= base64.StdEncoding.EncodedLen(numBytes) {
		numBytes++
	}
	b := make([]byte, numBytes)
	_, err := rand.Read(b)
	if err != nil {
		panic("could not generate cryptographically secure random string!")
	}

	return base64.StdEncoding.EncodeToString(b)
}

// NewOIDCSP creates a new internalOIDCProvider object from the given parameters
func NewOIDCSP(ctx context.Context, issuer, clientID, clientSecret, redirectURI string) (*InternalOIDCProvider, error) {
	o := &InternalOIDCProvider{
		Issuer:       issuer,
		ClientID:     clientID,
		ClientSecret: clientSecret,
		RedirectURI:  redirectURI,
	}
	var err error
	o.Provider, err = oidc.NewProvider(ctx, issuer)
	if err != nil {
		return nil, err
	}
	return o, nil
}

// GetOIDCLoginURL uses the given state to generate a login URL for the OIDC provider object
func (o *InternalOIDCProvider) GetOIDCLoginURL() (string, string, error) {
	if o == nil {
		return "", "", notConfigured
	}
	state := CryptoString(30)
	nonce := CryptoString(30)
	var err error
	conf := oauth2.Config{
		ClientID:     o.ClientID,
		ClientSecret: o.ClientSecret,
		RedirectURL:  o.RedirectURI,
		Endpoint:     o.Provider.Endpoint(),
		// "openid" is a required scope for OpenID Connect flows.
		// "profile" and "email" are necessary for using the email as an identifier
		Scopes: []string{oidc.ScopeOpenID, "profile", "email"},
	}

	si := stateInfoMap[state]
	si.Nonce = nonce
	stateInfoMap[state] = si

	url := conf.AuthCodeURL(state,
		oauth2.SetAuthURLParam("response_type", "code"),
		oauth2.SetAuthURLParam("nonce", nonce))
	return url, state, nil
}

// OIDCStateToEmail takes the state session created for the OIDC session
// and uses it discover the email of the user who obtained the
// code (or verify that the code belongs to them). This is how
// Pachyderm currently implements authorization in a production cluster
func (o *InternalOIDCProvider) OIDCStateToEmail(ctx context.Context, state string) (string, error) {
	// lookup the token from the given state
	si, ok := stateInfoMap[state]
	if !ok {
		return "", fmt.Errorf("did not have a valid state")
	}
	oauthToken := si.Token

	// Use the token passed in as our token source
	ts := oauth2.StaticTokenSource(
		&oauth2.Token{
			AccessToken: oauthToken,
		},
	)
	userInfo, err := o.Provider.UserInfo(ctx, ts)
	if err != nil {
		return "", err
	}
	logrus.Infof("recovered user info with email: '%v'", userInfo.Email)
	return userInfo.Email, nil
}

func (a *apiServer) handleExchange(w http.ResponseWriter, req *http.Request) {
	ctx := context.Background()

	var err error
	sp := a.getOIDCSP()
	if sp == nil {
		http.Error(w, "OIDC has not been configured or was disabled", http.StatusConflict)
		return
	}
	conf := &oauth2.Config{
		ClientID:     sp.ClientID,
		ClientSecret: sp.ClientSecret,
		RedirectURL:  sp.RedirectURI,
		Scopes:       []string{"openid", "email", "profile"},
		Endpoint:     sp.Provider.Endpoint(),
	}

	code := req.URL.Query()["code"][0]
	state := req.URL.Query()["state"][0]

	logrus.Infof("session state and code are obtained")

	// Use the authorization code that is pushed to the redirect
	tok, err := conf.Exchange(ctx, code)
	if err != nil {
		logrus.Errorf("failed to exchange code: %v", err)
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	logrus.Info("exchanged OIDC code for token")

	var verifier = sp.Provider.Verifier(&oidc.Config{ClientID: conf.ClientID})
	// Extract the ID Token from OAuth2 token.
	rawIDToken, ok := tok.Extra("id_token").(string)
	if !ok {
		logrus.Errorf("missing id token")
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	logrus.Infof("raw ID Token: %v", rawIDToken)

	// Parse and verify ID Token payload.
	idToken, err := verifier.Verify(ctx, rawIDToken)
	if err != nil {
		logrus.Errorf("could not verify token: %v", err)
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	si := stateInfoMap[state]
	if idToken.Nonce != si.Nonce {
		logrus.Errorf("expected nonces to match, instead set nonce %v but got nonce %v", si.Nonce, idToken.Nonce)
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}
	logrus.Infof("nonce is %v", idToken.Nonce)

	si.Token = tok.AccessToken
	logrus.Infof("saving state with access token")
	stateInfoMap[state] = si

	// let the CLI know that we've successfully exchanged the code, and verified the token
	tokenChan <- tokenInfo{token: tok.AccessToken, err: err}
	// make sure the channel is cleaned up for future logins
	close(tokenChan)
	tokenChan = make(chan tokenInfo)

	fmt.Fprintf(w, "You are now logged in. Go back to the terminal to use Pachyderm!")
}

func (a *apiServer) serveOIDC() {
	// serve OIDC handler to exchange the auth code
	http.HandleFunc("/authorization-code/callback", a.handleExchange)
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%v", OidcPort), nil))
}
<commit_msg>Add comments to InternalOIDCProvider<commit_after>package server

import (
	"crypto/rand"
	"encoding/base64"
	goerr "errors"
	"fmt"
	"log"
	"net/http"

	"github.com/coreos/go-oidc"
	logrus "github.com/sirupsen/logrus"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
)

// various oidc invalid argument errors. Use 'goerror' instead of internal
// 'errors' library b/c stack trace isn't useful
var (
	notConfigured = goerr.New("OIDC ID provider configuration not found")
)

var tokenChan chan tokenInfo

type sessionInfo struct {
	Nonce string
	Token string
}

var stateInfoMap map[string]sessionInfo

func init() {
	tokenChan = make(chan tokenInfo)
	stateInfoMap = make(map[string]sessionInfo)
}

// InternalOIDCProvider contains information about the configured OIDC ID
// provider, as well as auth information identifying Pachyderm in the ID
// provider (ClientID and ClientSecret), which Pachyderm needs to perform
// authorization with it.
type InternalOIDCProvider struct {
	// Provider generates the ID provider login URL returned by GetOIDCLogin
	Provider *oidc.Provider

	// Issuer is the address of the OIDC ID provider (where we exchange
	// authorization codes for access tokens and get users' email addresses in
	// Authorize())
	Issuer string

	// ClientID is Pachyderm's identifier in the OIDC ID provider (generated by
	// the ID provider, and passed to Pachyder by the cluster administrator via
	// SetConfig)
	ClientID string

	// ClientSecret is a shared secret with the ID provided, for doing the
	// auth-code -> access-token exchange.
	ClientSecret string
	// RedirectURI is used by GetOIDCLogin to generate a login URL that redirects
	// users back to pachyderm (must be provided by the cluster administrator via
	// SetConfig, as only they know their network topology & Pachyderm's address
	// within it, and must be included in login URLs)
	RedirectURI string

	// States is an etcd collection containing the state information associated
	// with every in-progress authentication flow. /authorization-code/callback
	// places users' ID tokens in here when they authenticate successfully, and
	// Authenticate() retrieves those ID tokens, converts them to Pachyderm
	// tokens, and returns users' Pachyderm tokens back to them—all scoped to the
	// state token identifying the login flow
	States col.Collection
}

// CryptoString returns a cryptographically random, URL safe string with length at least n
func CryptoString(n int) string {
	var numBytes int
	for n >= base64.StdEncoding.EncodedLen(numBytes) {
		numBytes++
	}
	b := make([]byte, numBytes)
	_, err := rand.Read(b)
	if err != nil {
		panic("could not generate cryptographically secure random string!")
	}

	return base64.StdEncoding.EncodeToString(b)
}

// NewOIDCSP creates a new internalOIDCProvider object from the given parameters
func NewOIDCSP(ctx context.Context, issuer, clientID, clientSecret, redirectURI string) (*InternalOIDCProvider, error) {
	o := &InternalOIDCProvider{
		Issuer:       issuer,
		ClientID:     clientID,
		ClientSecret: clientSecret,
		RedirectURI:  redirectURI,
	}
	var err error
	o.Provider, err = oidc.NewProvider(ctx, issuer)
	if err != nil {
		return nil, err
	}
	return o, nil
}

// GetOIDCLoginURL uses the given state to generate a login URL for the OIDC provider object
func (o *InternalOIDCProvider) GetOIDCLoginURL() (string, string, error) {
	if o == nil {
		return "", "", notConfigured
	}
	state := CryptoString(30)
	nonce := CryptoString(30)
	var err error
	conf := oauth2.Config{
		ClientID:     o.ClientID,
		ClientSecret: o.ClientSecret,
		RedirectURL:  o.RedirectURI,
		Endpoint:     o.Provider.Endpoint(),
		// "openid" is a required scope for OpenID Connect flows.
		// "profile" and "email" are necessary for using the email as an identifier
		Scopes: []string{oidc.ScopeOpenID, "profile", "email"},
	}

	si := stateInfoMap[state]
	si.Nonce = nonce
	stateInfoMap[state] = si

	url := conf.AuthCodeURL(state,
		oauth2.SetAuthURLParam("response_type", "code"),
		oauth2.SetAuthURLParam("nonce", nonce))
	return url, state, nil
}

// OIDCStateToEmail takes the state session created for the OIDC session
// and uses it discover the email of the user who obtained the
// code (or verify that the code belongs to them). This is how
// Pachyderm currently implements authorization in a production cluster
func (o *InternalOIDCProvider) OIDCStateToEmail(ctx context.Context, state string) (string, error) {
	// lookup the token from the given state
	si, ok := stateInfoMap[state]
	if !ok {
		return "", fmt.Errorf("did not have a valid state")
	}
	oauthToken := si.Token

	// Use the token passed in as our token source
	ts := oauth2.StaticTokenSource(
		&oauth2.Token{
			AccessToken: oauthToken,
		},
	)
	userInfo, err := o.Provider.UserInfo(ctx, ts)
	if err != nil {
		return "", err
	}
	logrus.Infof("recovered user info with email: '%v'", userInfo.Email)
	return userInfo.Email, nil
}

func (a *apiServer) handleExchange(w http.ResponseWriter, req *http.Request) {
	ctx := context.Background()

	var err error
	sp := a.getOIDCSP()
	if sp == nil {
		http.Error(w, "OIDC has not been configured or was disabled", http.StatusConflict)
		return
	}
	conf := &oauth2.Config{
		ClientID:     sp.ClientID,
		ClientSecret: sp.ClientSecret,
		RedirectURL:  sp.RedirectURI,
		Scopes:       []string{"openid", "email", "profile"},
		Endpoint:     sp.Provider.Endpoint(),
	}

	code := req.URL.Query()["code"][0]
	state := req.URL.Query()["state"][0]

	logrus.Infof("session state and code are obtained")

	// Use the authorization code that is pushed to the redirect
	tok, err := conf.Exchange(ctx, code)
	if err != nil {
		logrus.Errorf("failed to exchange code: %v", err)
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	logrus.Info("exchanged OIDC code for token")

	var verifier = sp.Provider.Verifier(&oidc.Config{ClientID: conf.ClientID})
	// Extract the ID Token from OAuth2 token.
	rawIDToken, ok := tok.Extra("id_token").(string)
	if !ok {
		logrus.Errorf("missing id token")
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	logrus.Infof("raw ID Token: %v", rawIDToken)

	// Parse and verify ID Token payload.
	idToken, err := verifier.Verify(ctx, rawIDToken)
	if err != nil {
		logrus.Errorf("could not verify token: %v", err)
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	si := stateInfoMap[state]
	if idToken.Nonce != si.Nonce {
		logrus.Errorf("expected nonces to match, instead set nonce %v but got nonce %v", si.Nonce, idToken.Nonce)
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}
	logrus.Infof("nonce is %v", idToken.Nonce)

	si.Token = tok.AccessToken
	logrus.Infof("saving state with access token")
	stateInfoMap[state] = si

	// let the CLI know that we've successfully exchanged the code, and verified the token
	tokenChan <- tokenInfo{token: tok.AccessToken, err: err}
	// make sure the channel is cleaned up for future logins
	close(tokenChan)
	tokenChan = make(chan tokenInfo)

	fmt.Fprintf(w, "You are now logged in. Go back to the terminal to use Pachyderm!")
}

func (a *apiServer) serveOIDC() {
	// serve OIDC handler to exchange the auth code
	http.HandleFunc("/authorization-code/callback", a.handleExchange)
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%v", OidcPort), nil))
}
<|endoftext|>


<commit_before>package mocknet

import (
	"context"
	"sync"
	"testing"
	"time"

	"github.com/libp2p/go-libp2p-core/network"
	"github.com/libp2p/go-libp2p-core/peer"
	ma "github.com/multiformats/go-multiaddr"
)

func TestNotifications(t *testing.T) {
	const swarmSize = 5
	const timeout = 10 * time.Second

	mn, err := FullMeshLinked(swarmSize)
	if err != nil {
		t.Fatal(err)
	}
	defer mn.Close()

	// signup notifs
	nets := mn.Nets()
	notifiees := make(map[peer.ID]*netNotifiee, len(nets))
	for _, pn := range nets {
		defer pn.Close()

		n := newNetNotifiee(t, swarmSize)
		pn.Notify(n)
		notifiees[pn.LocalPeer()] = n
	}

	// connect all but self
	if err := mn.ConnectAllButSelf(); err != nil {
		t.Fatal(err)
	}

	// test everyone got the correct connection opened calls
	for _, s1 := range nets {
		n := notifiees[s1.LocalPeer()]
		notifs := make(map[peer.ID][]network.Conn)
		for _, s2 := range nets {
			if s2 == s1 {
				continue
			}

			// this feels a little sketchy, but its probably okay
			for len(s1.ConnsToPeer(s2.LocalPeer())) != len(notifs[s2.LocalPeer()]) {
				select {
				case c := <-n.connected:
					nfp := notifs[c.RemotePeer()]
					notifs[c.RemotePeer()] = append(nfp, c)
				case <-time.After(timeout):
					t.Fatal("timeout")
				}
			}
		}

		for p, cons := range notifs {
			expect := s1.ConnsToPeer(p)
			if len(expect) != len(cons) {
				t.Fatal("got different number of connections")
			}

			for _, c := range cons {
				var found bool
				for _, c2 := range expect {
					if c == c2 {
						found = true
						break
					}
				}

				if !found {
					t.Fatal("connection not found!")
				}
			}
		}
	}

	for _, s := range nets {
		s.SetStreamHandler(func(s network.Stream) {
			s.Close()
		})
	}

	for _, s := range nets {
		s.SetStreamHandler(func(s network.Stream) {
			s.Close()
		})
	}

	// Make sure we've received at last one stream per conn.
	for _, s := range nets {
		conns := s.Conns()
		for _, c := range conns {
			st1, err := c.NewStream(context.Background())
			if err != nil {
				t.Error(err)
				continue
			}
			t.Logf("%s %s <--%p--> %s %s", c.LocalPeer(), c.LocalMultiaddr(), st1, c.RemotePeer(), c.RemoteMultiaddr())
			st1.Close()
		}
	}

	// close conns
	for _, s1 := range nets {
		n1 := notifiees[s1.LocalPeer()]
		for _, c1 := range s1.Conns() {
			c2 := ConnComplement(c1)

			n2 := notifiees[c2.LocalPeer()]
			c1.Close()

			var c3, c4 network.Conn
			select {
			case c3 = <-n1.disconnected:
			case <-time.After(timeout):
				t.Fatal("timeout")
			}
			if c1 != c3 {
				t.Fatal("got incorrect conn", c1, c3)
			}

			select {
			case c4 = <-n2.disconnected:
			case <-time.After(timeout):
				t.Fatal("timeout")
			}
			if c2 != c4 {
				t.Fatal("got incorrect conn", c1, c2)
			}
		}
	}

	for _, n1 := range notifiees {
		// Avoid holding this lock while waiting, otherwise we can deadlock.
		streamStateCopy := map[network.Stream]chan struct{}{}
		n1.streamState.Lock()
		for str, ch := range n1.streamState.m {
			streamStateCopy[str] = ch
		}
		n1.streamState.Unlock()

		for str1, ch1 := range streamStateCopy {
			<-ch1
			str2 := StreamComplement(str1)
			n2 := notifiees[str1.Conn().RemotePeer()]

			n2.streamState.Lock()
			ch2 := n2.streamState.m[str2]
			n2.streamState.Unlock()

			<-ch2
		}
	}
}

type netNotifiee struct {
	t *testing.T

	listen       chan ma.Multiaddr
	listenClose  chan ma.Multiaddr
	connected    chan network.Conn
	disconnected chan network.Conn

	streamState struct {
		sync.Mutex
		m map[network.Stream]chan struct{}
	}
}

func newNetNotifiee(t *testing.T, buffer int) *netNotifiee {
	nn := &netNotifiee{
		t:            t,
		listen:       make(chan ma.Multiaddr, 1),
		listenClose:  make(chan ma.Multiaddr, 1),
		connected:    make(chan network.Conn, buffer*2),
		disconnected: make(chan network.Conn, buffer*2),
	}
	nn.streamState.m = make(map[network.Stream]chan struct{})
	return nn
}

func (nn *netNotifiee) Listen(n network.Network, a ma.Multiaddr) {
	nn.listen <- a
}
func (nn *netNotifiee) ListenClose(n network.Network, a ma.Multiaddr) {
	nn.listenClose <- a
}
func (nn *netNotifiee) Connected(n network.Network, v network.Conn) {
	nn.connected <- v
}
func (nn *netNotifiee) Disconnected(n network.Network, v network.Conn) {
	nn.disconnected <- v
}
func (nn *netNotifiee) OpenedStream(n network.Network, s network.Stream) {
	nn.streamState.Lock()
	defer nn.streamState.Unlock()
	_, ok := nn.streamState.m[s]
	if ok {
		nn.t.Error("duplicate stream open")
		return
	}
	nn.streamState.m[s] = make(chan struct{})
}
func (nn *netNotifiee) ClosedStream(n network.Network, s network.Stream) {
	nn.streamState.Lock()
	defer nn.streamState.Unlock()
	ch, ok := nn.streamState.m[s]
	if !ok {
		nn.t.Error("saw close event but no open event")
		return
	}
	select {
	case <-ch:
		nn.t.Error("duplicate close event")
	default:
		close(ch)
	}
}
<commit_msg>fix flaky TestNotifications test<commit_after>package mocknet

import (
	"context"
	"sync"
	"testing"
	"time"

	"github.com/libp2p/go-libp2p-core/network"
	"github.com/libp2p/go-libp2p-core/peer"
	ma "github.com/multiformats/go-multiaddr"

	"github.com/stretchr/testify/require"
)

func TestNotifications(t *testing.T) {
	const swarmSize = 5
	const timeout = 10 * time.Second

	mn, err := FullMeshLinked(swarmSize)
	if err != nil {
		t.Fatal(err)
	}
	defer mn.Close()

	// signup notifs
	nets := mn.Nets()
	notifiees := make(map[peer.ID]*netNotifiee, len(nets))
	for _, pn := range nets {
		defer pn.Close()

		n := newNetNotifiee(t, swarmSize)
		pn.Notify(n)
		notifiees[pn.LocalPeer()] = n
	}

	// connect all but self
	if err := mn.ConnectAllButSelf(); err != nil {
		t.Fatal(err)
	}

	// test everyone got the correct connection opened calls
	for _, s1 := range nets {
		n := notifiees[s1.LocalPeer()]
		notifs := make(map[peer.ID][]network.Conn)
		for _, s2 := range nets {
			if s2 == s1 {
				continue
			}

			// this feels a little sketchy, but its probably okay
			for len(s1.ConnsToPeer(s2.LocalPeer())) != len(notifs[s2.LocalPeer()]) {
				select {
				case c := <-n.connected:
					nfp := notifs[c.RemotePeer()]
					notifs[c.RemotePeer()] = append(nfp, c)
				case <-time.After(timeout):
					t.Fatal("timeout")
				}
			}
		}

		for p, cons := range notifs {
			expect := s1.ConnsToPeer(p)
			if len(expect) != len(cons) {
				t.Fatal("got different number of connections")
			}

			for _, c := range cons {
				var found bool
				for _, c2 := range expect {
					if c == c2 {
						found = true
						break
					}
				}

				if !found {
					t.Fatal("connection not found!")
				}
			}
		}
	}

	acceptedStream := make(chan struct{}, 1000)
	for _, s := range nets {
		s.SetStreamHandler(func(s network.Stream) {
			acceptedStream <- struct{}{}
			s.Close()
		})
	}

	// Make sure we've received at last one stream per conn.
	for _, s := range nets {
		conns := s.Conns()
		for _, c := range conns {
			st1, err := c.NewStream(context.Background())
			if err != nil {
				t.Error(err)
				continue
			}
			t.Logf("%s %s <--%p--> %s %s", c.LocalPeer(), c.LocalMultiaddr(), st1, c.RemotePeer(), c.RemoteMultiaddr())
			st1.Close()
		}
	}

	// close conns
	for _, s1 := range nets {
		n1 := notifiees[s1.LocalPeer()]
		for _, c1 := range s1.Conns() {
			c2 := ConnComplement(c1)

			n2 := notifiees[c2.LocalPeer()]
			c1.Close()

			var c3, c4 network.Conn
			select {
			case c3 = <-n1.disconnected:
			case <-time.After(timeout):
				t.Fatal("timeout")
			}
			if c1 != c3 {
				t.Fatal("got incorrect conn", c1, c3)
			}

			select {
			case c4 = <-n2.disconnected:
			case <-time.After(timeout):
				t.Fatal("timeout")
			}
			if c2 != c4 {
				t.Fatal("got incorrect conn", c1, c2)
			}
		}
	}

	for _, n1 := range notifiees {
		// Avoid holding this lock while waiting, otherwise we can deadlock.
		streamStateCopy := map[network.Stream]chan struct{}{}
		n1.streamState.Lock()
		for str, ch := range n1.streamState.m {
			streamStateCopy[str] = ch
		}
		n1.streamState.Unlock()

		for str1, ch1 := range streamStateCopy {
			<-ch1
			str2 := StreamComplement(str1)
			n2 := notifiees[str1.Conn().RemotePeer()]

			// make sure the OpenedStream notification was processed first
			var ch2 chan struct{}
			require.Eventually(t, func() bool {
				n2.streamState.Lock()
				defer n2.streamState.Unlock()
				ch, ok := n2.streamState.m[str2]
				if ok {
					ch2 = ch
				}
				return ok
			}, time.Second, 10*time.Millisecond)

			<-ch2
		}
	}
}

type netNotifiee struct {
	t *testing.T

	listen       chan ma.Multiaddr
	listenClose  chan ma.Multiaddr
	connected    chan network.Conn
	disconnected chan network.Conn

	streamState struct {
		sync.Mutex
		m map[network.Stream]chan struct{}
	}
}

func newNetNotifiee(t *testing.T, buffer int) *netNotifiee {
	nn := &netNotifiee{
		t:            t,
		listen:       make(chan ma.Multiaddr, 1),
		listenClose:  make(chan ma.Multiaddr, 1),
		connected:    make(chan network.Conn, buffer*2),
		disconnected: make(chan network.Conn, buffer*2),
	}
	nn.streamState.m = make(map[network.Stream]chan struct{})
	return nn
}

func (nn *netNotifiee) Listen(n network.Network, a ma.Multiaddr) {
	nn.listen <- a
}
func (nn *netNotifiee) ListenClose(n network.Network, a ma.Multiaddr) {
	nn.listenClose <- a
}
func (nn *netNotifiee) Connected(n network.Network, v network.Conn) {
	nn.connected <- v
}
func (nn *netNotifiee) Disconnected(n network.Network, v network.Conn) {
	nn.disconnected <- v
}
func (nn *netNotifiee) OpenedStream(n network.Network, s network.Stream) {
	nn.streamState.Lock()
	defer nn.streamState.Unlock()
	_, ok := nn.streamState.m[s]
	if ok {
		nn.t.Error("duplicate stream open")
		return
	}
	nn.streamState.m[s] = make(chan struct{})
}
func (nn *netNotifiee) ClosedStream(n network.Network, s network.Stream) {
	nn.streamState.Lock()
	defer nn.streamState.Unlock()
	ch, ok := nn.streamState.m[s]
	if !ok {
		nn.t.Error("saw close event but no open event")
		return
	}
	select {
	case <-ch:
		nn.t.Error("duplicate close event")
	default:
		close(ch)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"testing"
)

const SIZE = 5

func TestEmptyMap(t *testing.T) {
	gameMap := initGameMap(SIZE, SIZE)

	if !isMapEmpty(gameMap) {
		t.Error("GameMap must be empty, but it is:", gameMap)
	}

	gameMap.Update()
	if !isMapEmpty(gameMap) {
		t.Error("GameMap must be empty, but it is:", gameMap)
	}
}

func TestCellBecomeAlive(t *testing.T) {
	gameMap := initGameMap(SIZE, SIZE)

	gameMap.SetValue(2, 1, true)
	gameMap.SetValue(1, 3, true)
	gameMap.SetValue(3, 3, true)
	//00000
	//00100
	//00!00
	//01010
	//00000

	gameMap.Update()
	//00000
	//00000
	//00100
	//00000
	//00000
	if !gameMap.GetValue(2, 2) {
		t.Error("Cell 2,3 must be alive, but it doesn't:", gameMap)
	}
}

func TestSimpleOscillator(t *testing.T) {
	gameMap := initGameMap(SIZE, SIZE)

	gameMap.SetValue(1, 2, true)
	gameMap.SetValue(2, 2, true)
	gameMap.SetValue(3, 2, true)
	//00000
	//00000
	//01110
	//00000
	//00000

	gameMap.Update()
	//00000
	//00100
	//00100
	//00100
	//00000
	if !gameMap.GetValue(2, 1) ||
		!gameMap.GetValue(2, 2) ||
		!gameMap.GetValue(2, 3) {
		t.Error("Oscillator must be vertical, but it doesn't:", gameMap)
	}
}

func isMapEmpty(gameMap *GameMap) bool {
	for i := 0; i < len(gameMap.cellMap); i ++ {
		for j := 0; j < len(gameMap.cellMap[0]); j++ {
			cellAlive := gameMap.GetValue(i, j)
			if cellAlive {
				return false
			}
		}
	}
	return true
}
<commit_msg>Additional basic test<commit_after>package main

import (
	"testing"
)

const SIZE = 5

func TestEmptyMap(t *testing.T) {
	gameMap := initGameMap(SIZE, SIZE)

	if !isMapEmpty(gameMap) {
		t.Error("GameMap must be empty, but it is:", gameMap)
	}

	gameMap.Update()
	if !isMapEmpty(gameMap) {
		t.Error("GameMap must be empty, but it is:", gameMap)
	}
}

func TestCellBecomeAlive(t *testing.T) {
	gameMap := initGameMap(SIZE, SIZE)

	gameMap.SetValue(2, 1, true)
	gameMap.SetValue(1, 3, true)
	gameMap.SetValue(3, 3, true)
	//00000
	//00100
	//00!00
	//01010
	//00000

	gameMap.Update()
	//00000
	//00000
	//00100
	//00000
	//00000

	if !gameMap.GetValue(2, 2) {
		t.Error("Cell 2,3 must be alive, but it doesn't:", gameMap)
	}
}

func TestCellDiesNoNeighbours(t *testing.T) {
	gameMap := initGameMap(SIZE, SIZE)

	gameMap.SetValue(2, 2, true)
	//00000
	//00000
	//00100
	//00000
	//00000

	gameMap.Update()
	//00000
	//00000
	//00000
	//00000
	//00000

	if gameMap.GetValue(2, 2) {
		t.Error("Cell 2,3 must be dead, but it doesn't:", gameMap)
	}
}

func TestSimpleOscillator(t *testing.T) {
	gameMap := initGameMap(SIZE, SIZE)

	gameMap.SetValue(1, 2, true)
	gameMap.SetValue(2, 2, true)
	gameMap.SetValue(3, 2, true)
	//00000
	//00000
	//01110
	//00000
	//00000

	gameMap.Update()
	//00000
	//00100
	//00100
	//00100
	//00000

	if !gameMap.GetValue(2, 1) ||
		!gameMap.GetValue(2, 2) ||
		!gameMap.GetValue(2, 3) {
		t.Error("Oscillator must be vertical, but it doesn't:", gameMap)
	}
}

func isMapEmpty(gameMap *GameMap) bool {
	for i := 0; i < len(gameMap.cellMap); i ++ {
		for j := 0; j < len(gameMap.cellMap[0]); j++ {
			cellAlive := gameMap.GetValue(i, j)
			if cellAlive {
				return false
			}
		}
	}
	return true
}
<|endoftext|>


<commit_before>package waiter

import (
	"fmt"
	"log"
	"regexp"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/cloudformation"
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/resource"
)

const (
	// Default maximum amount of time to wait for a StackSetInstance to be Created
	StackSetInstanceCreatedDefaultTimeout = 30 * time.Minute

	// Default maximum amount of time to wait for a StackSetInstance to be Updated
	StackSetInstanceUpdatedDefaultTimeout = 30 * time.Minute

	// Default maximum amount of time to wait for a StackSetInstance to be Deleted
	StackSetInstanceDeletedDefaultTimeout = 30 * time.Minute

	stackSetOperationDelay = 5 * time.Second
)

const (
	// Default maximum amount of time to wait for a StackSet to be Updated
	StackSetUpdatedDefaultTimeout = 30 * time.Minute
)

func StackSetOperationSucceeded(conn *cloudformation.CloudFormation, stackSetName, operationID string, timeout time.Duration) error {
	stateConf := &resource.StateChangeConf{
		Pending: []string{cloudformation.StackSetOperationStatusRunning},
		Target:  []string{cloudformation.StackSetOperationStatusSucceeded},
		Refresh: StackSetOperationStatus(conn, stackSetName, operationID),
		Timeout: timeout,
		Delay:   stackSetOperationDelay,
	}

	log.Printf("[DEBUG] Waiting for CloudFormation StackSet (%s) operation: %s", stackSetName, operationID)
	_, err := stateConf.WaitForState()

	return err
}

const (
	// Default maximum amount of time to wait for a Stack to be Created
	StackCreatedDefaultTimeout = 30 * time.Minute

	stackCreatedMinTimeout = 1 * time.Second

	// Default maximum amount of time to wait for a Stack to be Updated
	StackUpdatedDefaultTimeout = 30 * time.Minute

	stackUpdatedMinTimeout = 5 * time.Second

	// Default maximum amount of time to wait for a Stack to be Deleted
	StackDeletedDefaultTimeout = 30 * time.Minute

	stackDeletedMinTimeout = 5 * time.Second
)

func StackCreated(conn *cloudformation.CloudFormation, stackName string, timeout time.Duration) (*cloudformation.Stack, error) {
	stateConf := resource.StateChangeConf{
		Pending: []string{
			cloudformation.StackStatusCreateInProgress,
			cloudformation.StackStatusDeleteInProgress,
			cloudformation.StackStatusRollbackInProgress,
		},
		Target: []string{
			cloudformation.StackStatusCreateComplete,
			cloudformation.StackStatusCreateFailed,
			cloudformation.StackStatusDeleteComplete,
			cloudformation.StackStatusDeleteFailed,
			cloudformation.StackStatusRollbackComplete,
			cloudformation.StackStatusRollbackFailed,
		},
		Timeout:    timeout,
		MinTimeout: stackCreatedMinTimeout,
		Delay:      10 * time.Second,
		Refresh:    StackStatus(conn, stackName),
	}

	outputRaw, err := stateConf.WaitForState()
	if err != nil {
		return nil, err
	}

	stack, ok := outputRaw.(*cloudformation.Stack)
	if !ok {
		return nil, err
	}

	lastStatus := aws.StringValue(stack.StackStatus)
	switch lastStatus {
	// This will be the case if either disable_rollback is false or on_failure is ROLLBACK
	case cloudformation.StackStatusRollbackComplete, cloudformation.StackStatusRollbackFailed:
		reasons, err := GetCloudFormationRollbackReasons(stackName, nil, conn)
		if err != nil {
			return stack, fmt.Errorf("failed to create CloudFormation stack, rollback requested (%s). Got an error reading failure information: %w", lastStatus, err)
		}
		return stack, fmt.Errorf("failed to create CloudFormation stack, rollback requested (%s): %q", lastStatus, reasons)

		// This will be the case if on_failure is DELETE
	case cloudformation.StackStatusDeleteComplete, cloudformation.StackStatusDeleteFailed:
		reasons, err := GetCloudFormationDeletionReasons(stackName, conn)
		if err != nil {
			return stack, fmt.Errorf("failed to create CloudFormation stack, delete requested (%s). Got an error reading failure information: %w", lastStatus, err)
		}

		return stack, fmt.Errorf("failed to create CloudFormation stack, delete requested (%s): %q", lastStatus, reasons)

		// This will be the case if either disable_rollback is true or on_failure is DO_NOTHING
	case cloudformation.StackStatusCreateFailed:
		reasons, err := GetCloudFormationFailures(stackName, conn)
		if err != nil {
			return stack, fmt.Errorf("failed to create CloudFormation stack (%s). Got an error reading failure information: %w", lastStatus, err)
		}
		return stack, fmt.Errorf("failed to create CloudFormation stack (%s): %q", lastStatus, reasons)
	}

	return stack, nil
}

func StackUpdated(conn *cloudformation.CloudFormation, stackName string, lastUpdatedTime *time.Time, timeout time.Duration) (*cloudformation.Stack, error) {
	stateConf := resource.StateChangeConf{
		Pending: []string{
			cloudformation.StackStatusUpdateCompleteCleanupInProgress,
			cloudformation.StackStatusUpdateInProgress,
			cloudformation.StackStatusUpdateRollbackInProgress,
			cloudformation.StackStatusUpdateRollbackCompleteCleanupInProgress,
		},
		Target: []string{
			cloudformation.StackStatusCreateComplete,
			cloudformation.StackStatusUpdateComplete,
			cloudformation.StackStatusUpdateRollbackComplete,
			cloudformation.StackStatusUpdateRollbackFailed,
		},
		Timeout:    timeout,
		MinTimeout: stackUpdatedMinTimeout,
		Delay:      10 * time.Second,
		Refresh:    StackStatus(conn, stackName),
	}

	outputRaw, err := stateConf.WaitForState()
	if err != nil {
		return nil, err
	}

	stack, ok := outputRaw.(*cloudformation.Stack)
	if !ok {
		return nil, err
	}

	lastStatus := aws.StringValue(stack.StackStatus)
	if lastStatus == cloudformation.StackStatusUpdateRollbackComplete || lastStatus == cloudformation.StackStatusUpdateRollbackFailed {
		reasons, err := GetCloudFormationRollbackReasons(stackName, lastUpdatedTime, conn)
		if err != nil {
			return stack, fmt.Errorf("failed to update CloudFormation stack (%s). Got an error reading failure information: %w", lastStatus, err)
		}

		return stack, fmt.Errorf("failed to update CloudFormation stack (%s): %q", lastStatus, reasons)
	}

	return stack, nil
}

func StackDeleted(conn *cloudformation.CloudFormation, stackName string, timeout time.Duration) (*cloudformation.Stack, error) {
	stateConf := resource.StateChangeConf{
		Pending: []string{
			cloudformation.StackStatusDeleteInProgress,
			cloudformation.StackStatusRollbackInProgress,
		},
		Target: []string{
			cloudformation.StackStatusDeleteComplete,
			cloudformation.StackStatusDeleteFailed,
		},
		Timeout:    timeout,
		MinTimeout: stackDeletedMinTimeout,
		Delay:      10 * time.Second,
		Refresh:    StackStatus(conn, stackName),
	}

	outputRaw, err := stateConf.WaitForState()
	if err != nil {
		return nil, err
	}

	stack, ok := outputRaw.(*cloudformation.Stack)
	if !ok {
		return nil, err
	}

	lastStatus := aws.StringValue(stack.StackStatus)
	if lastStatus == cloudformation.StackStatusDeleteFailed {
		reasons, err := GetCloudFormationFailures(stackName, conn)
		if err != nil {
			return stack, fmt.Errorf("Failed getting reasons for deletion failure: %w", err)
		}

		return stack, fmt.Errorf("%s: %q", lastStatus, reasons)
	}

	return stack, nil
}

func GetCloudFormationDeletionReasons(stackId string, conn *cloudformation.CloudFormation) ([]string, error) {
	var failures []string

	err := conn.DescribeStackEventsPages(&cloudformation.DescribeStackEventsInput{
		StackName: aws.String(stackId),
	}, func(page *cloudformation.DescribeStackEventsOutput, lastPage bool) bool {
		for _, e := range page.StackEvents {
			if cfStackEventIsFailure(e) || cfStackEventIsStackDeletion(e) {
				failures = append(failures, *e.ResourceStatusReason)
			}
		}
		return !lastPage
	})

	return failures, err
}

func GetCloudFormationRollbackReasons(stackId string, afterTime *time.Time, conn *cloudformation.CloudFormation) ([]string, error) {
	var failures []string

	err := conn.DescribeStackEventsPages(&cloudformation.DescribeStackEventsInput{
		StackName: aws.String(stackId),
	}, func(page *cloudformation.DescribeStackEventsOutput, lastPage bool) bool {
		for _, e := range page.StackEvents {
			if afterTime != nil && !e.Timestamp.After(*afterTime) {
				continue
			}

			if cfStackEventIsFailure(e) || cfStackEventIsRollback(e) {
				failures = append(failures, *e.ResourceStatusReason)
			}
		}
		return !lastPage
	})

	return failures, err
}

func GetCloudFormationFailures(stackId string, conn *cloudformation.CloudFormation) ([]string, error) {
	var failures []string

	err := conn.DescribeStackEventsPages(&cloudformation.DescribeStackEventsInput{
		StackName: aws.String(stackId),
	}, func(page *cloudformation.DescribeStackEventsOutput, lastPage bool) bool {
		for _, e := range page.StackEvents {
			if cfStackEventIsFailure(e) {
				failures = append(failures, *e.ResourceStatusReason)
			}
		}
		return !lastPage
	})

	return failures, err
}

func cfStackEventIsFailure(event *cloudformation.StackEvent) bool {
	failRe := regexp.MustCompile("_FAILED$")
	return failRe.MatchString(*event.ResourceStatus) && event.ResourceStatusReason != nil
}

func cfStackEventIsRollback(event *cloudformation.StackEvent) bool {
	rollbackRe := regexp.MustCompile("^ROLLBACK_")
	return rollbackRe.MatchString(*event.ResourceStatus) && event.ResourceStatusReason != nil
}

func cfStackEventIsStackDeletion(event *cloudformation.StackEvent) bool {
	return *event.ResourceStatus == "DELETE_IN_PROGRESS" &&
		*event.ResourceType == "AWS::CloudFormation::Stack" &&
		event.ResourceStatusReason != nil
}
<commit_msg>Corrects error message<commit_after>package waiter

import (
	"fmt"
	"log"
	"regexp"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/cloudformation"
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/resource"
)

const (
	// Default maximum amount of time to wait for a StackSetInstance to be Created
	StackSetInstanceCreatedDefaultTimeout = 30 * time.Minute

	// Default maximum amount of time to wait for a StackSetInstance to be Updated
	StackSetInstanceUpdatedDefaultTimeout = 30 * time.Minute

	// Default maximum amount of time to wait for a StackSetInstance to be Deleted
	StackSetInstanceDeletedDefaultTimeout = 30 * time.Minute

	stackSetOperationDelay = 5 * time.Second
)

const (
	// Default maximum amount of time to wait for a StackSet to be Updated
	StackSetUpdatedDefaultTimeout = 30 * time.Minute
)

func StackSetOperationSucceeded(conn *cloudformation.CloudFormation, stackSetName, operationID string, timeout time.Duration) error {
	stateConf := &resource.StateChangeConf{
		Pending: []string{cloudformation.StackSetOperationStatusRunning},
		Target:  []string{cloudformation.StackSetOperationStatusSucceeded},
		Refresh: StackSetOperationStatus(conn, stackSetName, operationID),
		Timeout: timeout,
		Delay:   stackSetOperationDelay,
	}

	log.Printf("[DEBUG] Waiting for CloudFormation StackSet (%s) operation: %s", stackSetName, operationID)
	_, err := stateConf.WaitForState()

	return err
}

const (
	// Default maximum amount of time to wait for a Stack to be Created
	StackCreatedDefaultTimeout = 30 * time.Minute

	stackCreatedMinTimeout = 1 * time.Second

	// Default maximum amount of time to wait for a Stack to be Updated
	StackUpdatedDefaultTimeout = 30 * time.Minute

	stackUpdatedMinTimeout = 5 * time.Second

	// Default maximum amount of time to wait for a Stack to be Deleted
	StackDeletedDefaultTimeout = 30 * time.Minute

	stackDeletedMinTimeout = 5 * time.Second
)

func StackCreated(conn *cloudformation.CloudFormation, stackName string, timeout time.Duration) (*cloudformation.Stack, error) {
	stateConf := resource.StateChangeConf{
		Pending: []string{
			cloudformation.StackStatusCreateInProgress,
			cloudformation.StackStatusDeleteInProgress,
			cloudformation.StackStatusRollbackInProgress,
		},
		Target: []string{
			cloudformation.StackStatusCreateComplete,
			cloudformation.StackStatusCreateFailed,
			cloudformation.StackStatusDeleteComplete,
			cloudformation.StackStatusDeleteFailed,
			cloudformation.StackStatusRollbackComplete,
			cloudformation.StackStatusRollbackFailed,
		},
		Timeout:    timeout,
		MinTimeout: stackCreatedMinTimeout,
		Delay:      10 * time.Second,
		Refresh:    StackStatus(conn, stackName),
	}

	outputRaw, err := stateConf.WaitForState()
	if err != nil {
		return nil, err
	}

	stack, ok := outputRaw.(*cloudformation.Stack)
	if !ok {
		return nil, err
	}

	lastStatus := aws.StringValue(stack.StackStatus)
	switch lastStatus {
	// This will be the case if either disable_rollback is false or on_failure is ROLLBACK
	case cloudformation.StackStatusRollbackComplete, cloudformation.StackStatusRollbackFailed:
		reasons, err := GetCloudFormationRollbackReasons(stackName, nil, conn)
		if err != nil {
			return stack, fmt.Errorf("failed to create CloudFormation stack, rollback requested (%s). Got an error reading failure information: %w", lastStatus, err)
		}
		return stack, fmt.Errorf("failed to create CloudFormation stack, rollback requested (%s): %q", lastStatus, reasons)

		// This will be the case if on_failure is DELETE
	case cloudformation.StackStatusDeleteComplete, cloudformation.StackStatusDeleteFailed:
		reasons, err := GetCloudFormationDeletionReasons(stackName, conn)
		if err != nil {
			return stack, fmt.Errorf("failed to create CloudFormation stack, delete requested (%s). Got an error reading failure information: %w", lastStatus, err)
		}

		return stack, fmt.Errorf("failed to create CloudFormation stack, delete requested (%s): %q", lastStatus, reasons)

		// This will be the case if either disable_rollback is true or on_failure is DO_NOTHING
	case cloudformation.StackStatusCreateFailed:
		reasons, err := GetCloudFormationFailures(stackName, conn)
		if err != nil {
			return stack, fmt.Errorf("failed to create CloudFormation stack (%s). Got an error reading failure information: %w", lastStatus, err)
		}
		return stack, fmt.Errorf("failed to create CloudFormation stack (%s): %q", lastStatus, reasons)
	}

	return stack, nil
}

func StackUpdated(conn *cloudformation.CloudFormation, stackName string, lastUpdatedTime *time.Time, timeout time.Duration) (*cloudformation.Stack, error) {
	stateConf := resource.StateChangeConf{
		Pending: []string{
			cloudformation.StackStatusUpdateCompleteCleanupInProgress,
			cloudformation.StackStatusUpdateInProgress,
			cloudformation.StackStatusUpdateRollbackInProgress,
			cloudformation.StackStatusUpdateRollbackCompleteCleanupInProgress,
		},
		Target: []string{
			cloudformation.StackStatusCreateComplete,
			cloudformation.StackStatusUpdateComplete,
			cloudformation.StackStatusUpdateRollbackComplete,
			cloudformation.StackStatusUpdateRollbackFailed,
		},
		Timeout:    timeout,
		MinTimeout: stackUpdatedMinTimeout,
		Delay:      10 * time.Second,
		Refresh:    StackStatus(conn, stackName),
	}

	outputRaw, err := stateConf.WaitForState()
	if err != nil {
		return nil, err
	}

	stack, ok := outputRaw.(*cloudformation.Stack)
	if !ok {
		return nil, err
	}

	lastStatus := aws.StringValue(stack.StackStatus)
	if lastStatus == cloudformation.StackStatusUpdateRollbackComplete || lastStatus == cloudformation.StackStatusUpdateRollbackFailed {
		reasons, err := GetCloudFormationRollbackReasons(stackName, lastUpdatedTime, conn)
		if err != nil {
			return stack, fmt.Errorf("failed to update CloudFormation stack (%s). Got an error reading failure information: %w", lastStatus, err)
		}

		return stack, fmt.Errorf("failed to update CloudFormation stack (%s): %q", lastStatus, reasons)
	}

	return stack, nil
}

func StackDeleted(conn *cloudformation.CloudFormation, stackName string, timeout time.Duration) (*cloudformation.Stack, error) {
	stateConf := resource.StateChangeConf{
		Pending: []string{
			cloudformation.StackStatusDeleteInProgress,
			cloudformation.StackStatusRollbackInProgress,
		},
		Target: []string{
			cloudformation.StackStatusDeleteComplete,
			cloudformation.StackStatusDeleteFailed,
		},
		Timeout:    timeout,
		MinTimeout: stackDeletedMinTimeout,
		Delay:      10 * time.Second,
		Refresh:    StackStatus(conn, stackName),
	}

	outputRaw, err := stateConf.WaitForState()
	if err != nil {
		return nil, err
	}

	stack, ok := outputRaw.(*cloudformation.Stack)
	if !ok {
		return nil, err
	}

	lastStatus := aws.StringValue(stack.StackStatus)
	if lastStatus == cloudformation.StackStatusDeleteFailed {
		reasons, err := GetCloudFormationFailures(stackName, conn)
		if err != nil {
			return stack, fmt.Errorf("failed to delete CloudFormation stack (%s). Got an error reading failure information: %w", lastStatus, err)
		}

		return stack, fmt.Errorf("failed to delete CloudFormation stack (%s): %q", lastStatus, reasons)
	}

	return stack, nil
}

func GetCloudFormationDeletionReasons(stackId string, conn *cloudformation.CloudFormation) ([]string, error) {
	var failures []string

	err := conn.DescribeStackEventsPages(&cloudformation.DescribeStackEventsInput{
		StackName: aws.String(stackId),
	}, func(page *cloudformation.DescribeStackEventsOutput, lastPage bool) bool {
		for _, e := range page.StackEvents {
			if cfStackEventIsFailure(e) || cfStackEventIsStackDeletion(e) {
				failures = append(failures, *e.ResourceStatusReason)
			}
		}
		return !lastPage
	})

	return failures, err
}

func GetCloudFormationRollbackReasons(stackId string, afterTime *time.Time, conn *cloudformation.CloudFormation) ([]string, error) {
	var failures []string

	err := conn.DescribeStackEventsPages(&cloudformation.DescribeStackEventsInput{
		StackName: aws.String(stackId),
	}, func(page *cloudformation.DescribeStackEventsOutput, lastPage bool) bool {
		for _, e := range page.StackEvents {
			if afterTime != nil && !e.Timestamp.After(*afterTime) {
				continue
			}

			if cfStackEventIsFailure(e) || cfStackEventIsRollback(e) {
				failures = append(failures, *e.ResourceStatusReason)
			}
		}
		return !lastPage
	})

	return failures, err
}

func GetCloudFormationFailures(stackId string, conn *cloudformation.CloudFormation) ([]string, error) {
	var failures []string

	err := conn.DescribeStackEventsPages(&cloudformation.DescribeStackEventsInput{
		StackName: aws.String(stackId),
	}, func(page *cloudformation.DescribeStackEventsOutput, lastPage bool) bool {
		for _, e := range page.StackEvents {
			if cfStackEventIsFailure(e) {
				failures = append(failures, *e.ResourceStatusReason)
			}
		}
		return !lastPage
	})

	return failures, err
}

func cfStackEventIsFailure(event *cloudformation.StackEvent) bool {
	failRe := regexp.MustCompile("_FAILED$")
	return failRe.MatchString(*event.ResourceStatus) && event.ResourceStatusReason != nil
}

func cfStackEventIsRollback(event *cloudformation.StackEvent) bool {
	rollbackRe := regexp.MustCompile("^ROLLBACK_")
	return rollbackRe.MatchString(*event.ResourceStatus) && event.ResourceStatusReason != nil
}

func cfStackEventIsStackDeletion(event *cloudformation.StackEvent) bool {
	return *event.ResourceStatus == "DELETE_IN_PROGRESS" &&
		*event.ResourceType == "AWS::CloudFormation::Stack" &&
		event.ResourceStatusReason != nil
}
<|endoftext|>


<commit_before>package main

import (
    "fmt"
    //"html"
    "log"
    "net/http"
	"io/ioutil"
	"strings"
)

type viewHandler struct{}

var (
	head string = `<!DOCTYPE html>

<title>Page Title</title>

This is a Heading
`
	tail string =	`</body></html>`
)

func main() {
	http.HandleFunc("/List", List)
	http.HandleFunc("/Watch/", Watch)
	http.Handle("/", http.FileServer(http.Dir("./Vids/")))
	//http.Handle("/", new(viewHandler))

    //http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
    //    fmt.Fprintf(w, "Hello, %q", html.EscapeString(r.URL.Path))
    //})

    log.Fatal(http.ListenAndServe(":8080", nil))

}

func List(w http.ResponseWriter, r *http.Request){
	fileText:=""
	files, err := ioutil.ReadDir("./Vids")
	if(err != nil){
		fileText = `<p>No Files Found</p>`
	}
	for _, file := range files {
		if strings.Contains(strings.ToLower(file.Name()), "mp4"){
			fileText += (""+file.Name()+"\n
\n")
		}else{
			fileText += ("
" + file.Name() + "
\n
\n")
		}

	}
	fmt.Fprintf(w, head + fileText + tail)
}

func Watch(w http.ResponseWriter, r *http.Request){
	fileToWatch := r.URL.Path[len("/view/"):]
	fileText := ("<video width=\"400\" preload=\"none\" controls><source src=\""+fileToWatch+"\" type=\"video/mp4\">Your browser does not support HTML5 video.</video>")
	fmt.Fprintf(w, head + fileText + tail)
}
<commit_msg>Added script to improve video start performace<commit_after>package main

import (
    "fmt"
    //"html"
    "log"
    "net/http"
	"io/ioutil"
	"strings"
)

type viewHandler struct{}

var (
	head string = `<!DOCTYPE html>

<title>Page Title</title>

This is a Heading
`

	tail string =	`</body></html>`

	script string = ``
)

func main() {
	http.HandleFunc("/List", List)
	http.HandleFunc("/Watch/", Watch)
	http.Handle("/", http.FileServer(http.Dir("./Vids/")))
	//http.Handle("/", new(viewHandler))

    //http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
    //    fmt.Fprintf(w, "Hello, %q", html.EscapeString(r.URL.Path))
    //})

    log.Fatal(http.ListenAndServe(":8080", nil))

}

func List(w http.ResponseWriter, r *http.Request){
	fileText:=""
	files, err := ioutil.ReadDir("./Vids")
	if(err != nil){
		fileText = `<p>No Files Found</p>`
	}
	for _, file := range files {
		if strings.Contains(strings.ToLower(file.Name()), "mp4"){
			fileText += (""+file.Name()+"\n
\n")
		}else{
			fileText += ("
" + file.Name() + "
\n
\n")
		}

	}
	fmt.Fprintf(w, head + fileText + tail)
}

func Watch(w http.ResponseWriter, r *http.Request){
	fileToWatch := r.URL.Path[len("/view/"):]
	fileText := ("<video width=\"400\" preload=\"none\" controls><source src=\""+fileToWatch+"\" type=\"video/mp4\">Your browser does not support HTML5 video.</video>")
	fmt.Fprintf(w, head + fileText + script + tail)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"log"
	"os"
)

func main() {
    flag.Parse()

    // receives arguments from command line
	args := flag.Args()

    // check if there are exactly two arguments (source and destination of hard link)
	if len(args) != 2 {
        log.Println(args)
		log.Fatal("Link script must receive exactly two arguments: source and destination")
	}

	src := args[0]
	dst := args[1]

    // make a direct call to the syscall of current OS
	if err := os.Link(src, dst); err != nil {
        // if there is an error, prints error message and exits
        log.Println("An error occurred while trying to create requested link")
		log.Fatal(err)
	}

	log.Println("Successfully created requested link")
}
<commit_msg>Added flag -u, to allow hard link removing<commit_after>package main

import (
	"flag"
	"log"
	"os"
)

func main() {
    unln := flag.Bool("u", false, "Unlinks the specified destination path")

    flag.Parse()

    // receives arguments from command line
	args := flag.Args()

    if *unln { // if -u flag was provided, do unlink destination
        // check if there is exactly one arguments (destination of hard link)
        if len(args) != 1 {
            log.Println(args)
            log.Fatal("Link script must receive exactly one argument when unlinking: destination")
        }

        dst := args[0] // destination

        // make a direct call to the syscall of current OS
        if err := os.Remove(dst); err != nil {
            // if there is an error, prints error message and exits
            log.Println("An error occurred while trying to remove specified link")
            log.Fatal(err)
        }

        log.Println("Successfully removed specified link")
    } else { // otherwise, do link source to destination
        // check if there are exactly two arguments (source and destination of hard link)
        if len(args) != 2 {
            log.Println(args)
            log.Fatal("Link script must receive exactly two arguments when linking: source and destination")
        }

        src := args[0] // source
        dst := args[1] // destination

        // make a direct call to the syscall of current OS
        if err := os.Link(src, dst); err != nil {
            // if there is an error, prints error message and exits
            log.Println("An error occurred while trying to create specified link")
            log.Fatal(err)
        }

        log.Println("Successfully created specified link")
    }

}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"

	"flag"
	"github.com/br0r/hydra/config"
)

const BASE_DIR string = ".hydra"
const help string = `Usage:
hydra [OPTIONS] COMMAND

Commands:
  init - Create hydra project [--clean]
  start - Start hydra servers
  kill - Kills started servers
  ls - Show started servers
`

func initialize(c config.Config, clean bool) {
	if clean {
		_, e := exec.Command("rm", "-rf", ".hydra").Output()
		if e != nil {
			log.Fatal(e)
		}
	}

	exec.Command("mkdir", ".hydra").Output()

	for i := 0; i < len(c.Services); i += 1 {
		service := c.Services[i]
		p := service.Path
		name := service.Name
		base := path.Join(".hydra/", name)
		if _, err := os.Stat(base); os.IsNotExist(err) {
			fmt.Println("Creating", name)

			var e error
			if strings.Index(p, "git@") == 0 {
				e = exec.Command("git", "clone", p, base).Run()
			} else {
				e = exec.Command("cp", "-R", p, base).Run()
			}

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

			git := path.Join(base, ".git")
			exec.Command("rm", "-rf", git).Run()

			if service.Config.Dest != "" && service.Config.Src != "" {
				u := path.Join(base, service.Config.Dest)
				cf := service.Config.Src

				err := exec.Command("cp", cf, u).Run()
				if err != nil {
					log.Fatal(err)
				}
			}

			if service.Install != "" {
				args := strings.Split(service.Install, " ")
				cmd := exec.Command(args[0], args[1:]...)

				cmd.Dir = base
				e := cmd.Run()
				if e != nil {
					log.Fatal(e)
				}
			}
		}
	}
}

func start(conf config.Config) {
	fmt.Println("Starting")
	for i := 0; i < len(conf.Services); i += 1 {
		service := conf.Services[i]
		name := service.Name

		cmdRunDir := path.Join(".hydra", name)
		pid_file := path.Join(cmdRunDir, "pid")
		if _, err := os.Stat(pid_file); !os.IsNotExist(err) {
			fmt.Printf("%s is already started, run hydra kill if you want to restart\n", name)
			continue
		}

		env := os.Environ()
		for k, v := range service.Env {
			env = append(env, fmt.Sprintf("%s=%s", k, v))
		}

		fmt.Println("Starting", name)
		args := strings.Split(conf.Services[i].Start, " ")
		cmd := exec.Command(args[0], args[1:]...)
		cmd.Env = env
		cmd.Dir = cmdRunDir
		cmd.Stderr = os.Stderr

		e := cmd.Start()
		pid := []byte(strconv.Itoa(cmd.Process.Pid))
		ioutil.WriteFile(pid_file, pid, 0440)

		if e != nil {
			os.Remove("pid")
			log.Fatal("Error with", name, e)
		}
	}
}

func kill(conf config.Config) {
	for i := 0; i < len(conf.Services); i += 1 {
		service := conf.Services[i]
		name := service.Name
		pid_file := path.Join(".hydra", name, "pid")
		if _, err := os.Stat(pid_file); !os.IsNotExist(err) {
			bytes, e := ioutil.ReadFile(pid_file)
			if e != nil {
				log.Fatalf("Error when reading file %v", e)
			}
			pid := string(bytes)
			exec.Command("kill", pid).Run()
			if e == nil {
				fmt.Printf("Killed %s\n", name)
			}
			os.Remove(pid_file)
		}
	}
}

func ls(conf config.Config) {
	for _, service := range conf.Services {
		name := service.Name
		pid_path := path.Join(BASE_DIR, name, "pid")
		if _, err := os.Stat(pid_path); !os.IsNotExist(err) {
			bytes, e := ioutil.ReadFile(pid_path)
			if e != nil {
				log.Fatal(e)
			}
			pid := string(bytes)
			fmt.Printf("%s running on pid: %s\n", name, pid)
		}
	}
}

func main() {
	var c config.Config = config.ReadConfig()
	var clean = flag.Bool("clean", false, "If we want a clean run")
	flag.Parse()

	if cmd := flag.Arg(0); cmd != "" {
		switch cmd {
		case "init":
			initialize(c, *clean)
		case "start":
			start(c)
		case "kill":
			kill(c)
		case "ls":
			ls(c)
		default:
			fmt.Println(help)
		}
	} else {
		fmt.Println(help)
	}
}
<commit_msg>started with logs<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"

	"flag"
	"github.com/br0r/hydra/config"
)

const BASE_DIR string = ".hydra"
const help string = `Usage:
hydra [OPTIONS] COMMAND

Commands:
  init - Create hydra project [--clean]
  start - Start hydra servers
  stop - Stops started servers
  ls - Show started servers
  logs [name] - Show logs for servers, or for specific server given by name
`

func initialize(c config.Config, clean bool) {
	if clean {
		_, e := exec.Command("rm", "-rf", ".hydra").Output()
		if e != nil {
			log.Fatal(e)
		}
	}

	exec.Command("mkdir", ".hydra").Output()

	for i := 0; i < len(c.Services); i += 1 {
		service := c.Services[i]
		p := service.Path
		name := service.Name
		base := path.Join(".hydra/", name)
		if _, err := os.Stat(base); os.IsNotExist(err) {
			fmt.Println("Creating", name)

			var e error
			if strings.Index(p, "git@") == 0 {
				e = exec.Command("git", "clone", p, base).Run()
			} else {
				e = exec.Command("cp", "-R", p, base).Run()
			}

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

			git := path.Join(base, ".git")
			exec.Command("rm", "-rf", git).Run()

			if service.Config.Dest != "" && service.Config.Src != "" {
				u := path.Join(base, service.Config.Dest)
				cf := service.Config.Src

				err := exec.Command("cp", cf, u).Run()
				if err != nil {
					log.Fatal(err)
				}
			}

			if service.Install != "" {
				args := strings.Split(service.Install, " ")
				cmd := exec.Command(args[0], args[1:]...)

				cmd.Dir = base
				e := cmd.Run()
				if e != nil {
					log.Fatal(e)
				}
			}
		}
	}
}

func start(conf config.Config) {
	fmt.Println("Starting")
	for i := 0; i < len(conf.Services); i += 1 {
		service := conf.Services[i]
		name := service.Name

		cmdRunDir := path.Join(".hydra", name)
		pid_file := path.Join(cmdRunDir, "pid")
		if _, err := os.Stat(pid_file); !os.IsNotExist(err) {
			fmt.Printf("%s is already started, run hydra kill if you want to restart\n", name)
			continue
		}

		env := os.Environ()
		for k, v := range service.Env {
			env = append(env, fmt.Sprintf("%s=%s", k, v))
		}

		fmt.Println("Starting", name)
		args := strings.Split(conf.Services[i].Start, " ")
		cmd := exec.Command(args[0], args[1:]...)
		cmd.Env = env
		cmd.Dir = cmdRunDir
		cmd.Stderr = os.Stderr

		e := cmd.Start()
		pid := []byte(strconv.Itoa(cmd.Process.Pid))
		ioutil.WriteFile(pid_file, pid, 0440)

		if e != nil {
			os.Remove("pid")
			log.Fatal("Error with", name, e)
		}
	}
}

func kill(conf config.Config) {
	for i := 0; i < len(conf.Services); i += 1 {
		service := conf.Services[i]
		name := service.Name
		pid_file := path.Join(".hydra", name, "pid")
		if _, err := os.Stat(pid_file); !os.IsNotExist(err) {
			bytes, e := ioutil.ReadFile(pid_file)
			if e != nil {
				log.Fatalf("Error when reading file %v", e)
			}
			pid := string(bytes)
			exec.Command("kill", pid).Run()
			if e == nil {
				fmt.Printf("Killed %s\n", name)
			}
			os.Remove(pid_file)
		}
	}
}

func ls(conf config.Config) {
	for _, service := range conf.Services {
		name := service.Name
		pid_path := path.Join(BASE_DIR, name, "pid")
		if _, err := os.Stat(pid_path); !os.IsNotExist(err) {
			bytes, e := ioutil.ReadFile(pid_path)
			if e != nil {
				log.Fatal(e)
			}
			pid := string(bytes)
			fmt.Printf("%s running on pid: %s\n", name, pid)
		}
	}
}

func logs(conf config.Config, string []servers) {
	for _, service := range conf.Services {
		if len(servers) > 0 {
			err := false
			for _, name := range servers {
				if name != service.Name {
					err = true
					break
				}
			}
			if err == true {
				continue
			}
		}
	}
}

func main() {
	var c config.Config = config.ReadConfig()
	var clean = flag.Bool("clean", false, "If we want a clean run")
	flag.Parse()

	if cmd := flag.Arg(0); cmd != "" {
		switch cmd {
		case "init":
			initialize(c, *clean)
		case "start":
			start(c)
		case "stop":
			kill(c)
		case "ls":
			ls(c)
		case "logs":
			logs(c, flag.Args)
		default:
			fmt.Println(help)
		}
	} else {
		fmt.Println(help)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"strings"
)

var separator = os.Getenv("BRONZE_SEPARATOR")
var thinSeparator = os.Getenv("BRONZE_THIN_SEPARATOR")
var iconMode = os.Getenv("BRONZE_ICONS")
var icons = map[string]string{}

func init() {
	// set defaults
	if separator == "" {
		separator = "\ue0b0" // Powerline
	}
	if thinSeparator == "" {
		thinSeparator = "\ue0b1" // Powerline
	}
	if iconMode == "" {
		iconMode = "nerd"
	}

	switch iconMode {
	case "nerd":
		initIcon("apple", "\uf179")      // Font Awesome; apple
		initIcon("arch", "\uf300")       // Font Linux
		initIcon("centOS", "\uf301")     // Font Linux
		initIcon("debian", "\uf302")     // Font Linux
		initIcon("fedora", "\uf303")     // Font Linux
		initIcon("mint", "\uf304")       // Font Linux
		initIcon("SUSE", "\uf308")       // Font Linux
		initIcon("ubuntu", "\uf30c")     // Font Linux
		initIcon("elementary", "\uf311") // Font Linux
		initIcon("linux", "\uf17c")      // Font Awesome; linux
		initIcon("bsd", "\uf10f")        // Font Linux
		initIcon("root", "\ue00a")       // Pomicons; lightning
		initIcon("readonly", "\uf023")   // Font Awesome; lock
		initIcon("failed", "\ue009")     // Pomicons; exclamation
		initIcon("job", "\ue615")        // Seti UI
		initIcon("package", "\uf187")    // Font Awesome; archive
		initIcon("home", "\uf015")       // Font Awesome; home
		initIcon("github", "\uf09b")     // Font Awesome; github
		initIcon("gitlab", "\uf296")     // Font Awesome; gitlab
		initIcon("bitbucket", "\uf171")  // Font Awesome; bitbucket
		initIcon("git", "\ue0a0")        // Powerline
		initIcon("stash", "\uf01c")      // Font Awesome; inbox
		initIcon("ahead", "\uf148")      // Font Awesome; level-up
		initIcon("behind", "\uf149")     // Font Awesome; level-down
		initIcon("modified", "\u25cf")   // Unicode
		initIcon("staged", "\u271a")     // Unicode
	case "unicode":
		// TODO: test if it's possible to use \uf8ff on an Apple machine
		initIcon("apple", "\U0001f34e") // Emoji; red apple
		initIcon("arch", "")
		initIcon("centOS", "")
		initIcon("debian", "")
		initIcon("fedora", "")
		initIcon("mint", "")
		initIcon("SUSE", "")
		initIcon("ubuntu", "")
		initIcon("elementary", "")
		initIcon("linux", "\U0001f427") // Emoji; penguin
		initIcon("bsd", "")
		initIcon("root", "\u26a1")         // Emoji; high voltage
		initIcon("readonly", "\U0001f512") // Emoji; locked
		initIcon("failed", "\u2757")       // Emoji; exclamation mark
		initIcon("job", "\u2699")          // Emoji; gear
		initIcon("package", "\U0001f4e6")  // Emoji; package
		initIcon("home", "\U0001f3e0")     // Emoji; house
		initIcon("github", "")
		initIcon("gitlab", "")
		initIcon("bitbucket", "")
		initIcon("git", "")
		initIcon("stash", "\U0001f4e5") // Emoji; inbox tray
		initIcon("ahead", "\u2191")     // Unicode
		initIcon("behind", "\u2193")    // Unicode
		initIcon("modified", "\u25cf")  // Unicode
		initIcon("staged", "\u271a")    // Unicode
	case "ascii":
		initIcon("apple", "")
		initIcon("arch", "")
		initIcon("centOS", "")
		initIcon("debian", "")
		initIcon("fedora", "")
		initIcon("mint", "")
		initIcon("SUSE", "")
		initIcon("ubuntu", "")
		initIcon("elementary", "")
		initIcon("linux", "")
		initIcon("bsd", "")
		initIcon("root", "#")
		initIcon("readonly", "@")
		initIcon("failed", "!")
		initIcon("job", "&")
		initIcon("package", "")
		initIcon("home", "~")
		initIcon("github", "")
		initIcon("gitlab", "")
		initIcon("bitbucket", "")
		initIcon("git", "")
		initIcon("stash", "#")
		initIcon("ahead", ">")
		initIcon("behind", "<")
		initIcon("modified", "*")
		initIcon("staged", "+")
	}
}

func initIcon(name, defaultValue string) {
	env := os.Getenv("BRONZE_ICON_" + strings.ToUpper(name))
	if env == "" {
		icons[name] = defaultValue
	} else {
		icons[name] = env
	}
}
<commit_msg>fixed bsd unicode codepoint<commit_after>package main

import (
	"os"
	"strings"
)

var separator = os.Getenv("BRONZE_SEPARATOR")
var thinSeparator = os.Getenv("BRONZE_THIN_SEPARATOR")
var iconMode = os.Getenv("BRONZE_ICONS")
var icons = map[string]string{}

func init() {
	// set defaults
	if separator == "" {
		separator = "\ue0b0" // Powerline
	}
	if thinSeparator == "" {
		thinSeparator = "\ue0b1" // Powerline
	}
	if iconMode == "" {
		iconMode = "nerd"
	}

	switch iconMode {
	case "nerd":
		initIcon("apple", "\uf179")      // Font Awesome; apple
		initIcon("arch", "\uf300")       // Font Linux
		initIcon("centOS", "\uf301")     // Font Linux
		initIcon("debian", "\uf302")     // Font Linux
		initIcon("fedora", "\uf303")     // Font Linux
		initIcon("mint", "\uf304")       // Font Linux
		initIcon("SUSE", "\uf308")       // Font Linux
		initIcon("ubuntu", "\uf30c")     // Font Linux
		initIcon("elementary", "\uf311") // Font Linux
		initIcon("linux", "\uf17c")      // Font Awesome; linux
		initIcon("bsd", "\uf30e")        // Font Linux
		initIcon("root", "\ue00a")       // Pomicons; lightning
		initIcon("readonly", "\uf023")   // Font Awesome; lock
		initIcon("failed", "\ue009")     // Pomicons; exclamation
		initIcon("job", "\ue615")        // Seti UI
		initIcon("package", "\uf187")    // Font Awesome; archive
		initIcon("home", "\uf015")       // Font Awesome; home
		initIcon("github", "\uf09b")     // Font Awesome; github
		initIcon("gitlab", "\uf296")     // Font Awesome; gitlab
		initIcon("bitbucket", "\uf171")  // Font Awesome; bitbucket
		initIcon("git", "\ue0a0")        // Powerline
		initIcon("stash", "\uf01c")      // Font Awesome; inbox
		initIcon("ahead", "\uf148")      // Font Awesome; level-up
		initIcon("behind", "\uf149")     // Font Awesome; level-down
		initIcon("modified", "\u25cf")   // Unicode
		initIcon("staged", "\u271a")     // Unicode
	case "unicode":
		// TODO: test if it's possible to use \uf8ff on an Apple machine
		initIcon("apple", "\U0001f34e") // Emoji; red apple
		initIcon("arch", "")
		initIcon("centOS", "")
		initIcon("debian", "")
		initIcon("fedora", "")
		initIcon("mint", "")
		initIcon("SUSE", "")
		initIcon("ubuntu", "")
		initIcon("elementary", "")
		initIcon("linux", "\U0001f427") // Emoji; penguin
		initIcon("bsd", "")
		initIcon("root", "\u26a1")         // Emoji; high voltage
		initIcon("readonly", "\U0001f512") // Emoji; locked
		initIcon("failed", "\u2757")       // Emoji; exclamation mark
		initIcon("job", "\u2699")          // Emoji; gear
		initIcon("package", "\U0001f4e6")  // Emoji; package
		initIcon("home", "\U0001f3e0")     // Emoji; house
		initIcon("github", "")
		initIcon("gitlab", "")
		initIcon("bitbucket", "")
		initIcon("git", "")
		initIcon("stash", "\U0001f4e5") // Emoji; inbox tray
		initIcon("ahead", "\u2191")     // Unicode
		initIcon("behind", "\u2193")    // Unicode
		initIcon("modified", "\u25cf")  // Unicode
		initIcon("staged", "\u271a")    // Unicode
	case "ascii":
		initIcon("apple", "")
		initIcon("arch", "")
		initIcon("centOS", "")
		initIcon("debian", "")
		initIcon("fedora", "")
		initIcon("mint", "")
		initIcon("SUSE", "")
		initIcon("ubuntu", "")
		initIcon("elementary", "")
		initIcon("linux", "")
		initIcon("bsd", "")
		initIcon("root", "#")
		initIcon("readonly", "@")
		initIcon("failed", "!")
		initIcon("job", "&")
		initIcon("package", "")
		initIcon("home", "~")
		initIcon("github", "")
		initIcon("gitlab", "")
		initIcon("bitbucket", "")
		initIcon("git", "")
		initIcon("stash", "#")
		initIcon("ahead", ">")
		initIcon("behind", "<")
		initIcon("modified", "*")
		initIcon("staged", "+")
	}
}

func initIcon(name, defaultValue string) {
	env := os.Getenv("BRONZE_ICON_" + strings.ToUpper(name))
	if env == "" {
		icons[name] = defaultValue
	} else {
		icons[name] = env
	}
}
<|endoftext|>


<commit_before>package mdocker

import (
	"archive/tar"
	"bufio"
	"compress/gzip"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/fsouza/go-dockerclient"
	"github.com/mistifyio/mistify-agent/rpc"
)

// ListImages retrieves a list of Docker images
func (md *MDocker) ListImages(h *http.Request, request *rpc.ImageRequest, response *rpc.ImageResponse) error {
	opts := docker.ListImagesOptions{}

	apiImages, err := md.client.ListImages(opts)
	if err != nil {
		return err
	}
	images := make([]*rpc.Image, 0, len(apiImages))
	for _, ai := range apiImages {
		id, _ := docker.ParseRepositoryTag(ai.RepoTags[0])
		images = append(images, &rpc.Image{
			Id:   id,
			Type: "container",
			Size: uint64(ai.Size) / 1024 / 1024,
		})
	}

	response.Images = images
	return nil
}

// GetImage retrieves information about a specific Docker image
func (md *MDocker) GetImage(h *http.Request, request *rpc.ImageRequest, response *rpc.ImageResponse) error {
	image, err := md.client.InspectImage(request.Id)
	if err != nil {
		return err
	}

	response.Images = []*rpc.Image{
		&rpc.Image{
			Id:   request.Id,
			Type: "container",
			Size: uint64(image.Size) / 1024 / 1024,
		},
	}
	return nil
}

// LoadImage downloads a new container image from the image service and
// imports it into Docker
func (md *MDocker) LoadImage(h *http.Request, request *rpc.ImageRequest, response *rpc.ImageResponse) error {
	name := request.Id

	// Check if we already have the image to avoid unnecessary pulling
	image, err := md.client.InspectImage(name)
	if err != nil && err != docker.ErrNoSuchImage {
		return err
	}

	if image == nil {
		// Docker Import lets the image get renamed, but it strips metadata
		// (which includes any CMD that had been set). Docker Load doesn't let
		// the image get renamed and doesn't return the image id or name after
		// loading, but does preserve metadata. Since a rename to use the
		// image-service assigned id and metadata with CMD are both necessary,
		// the only way forward is to update the repositories file inside and
		// then load it.
		source := fmt.Sprintf("http://%s/images/%s/download", md.imageService, request.Id)
		resp, err := http.Get(source)
		if err != nil {
			return err
		}
		defer resp.Body.Close()

		if resp.StatusCode != http.StatusOK {
			return ErrorHTTPCode{
				Expected: http.StatusOK,
				Code:     resp.StatusCode,
				Source:   source,
			}
		}

		// Use a response buffer so the first few bytes can be peeked at for
		// file type detection. Uncompress the image if it is gzipped
		responseBuffer := bufio.NewReader(resp.Body)
		var imageReader io.Reader = responseBuffer
		filetypeBytes, err := responseBuffer.Peek(512)
		if err != nil {
			return err
		}
		if http.DetectContentType(filetypeBytes) == "application/x-gzip" {
			gzipReader, err := gzip.NewReader(responseBuffer)
			if err != nil {
				return err
			}
			defer gzipReader.Close()
			imageReader = gzipReader
		}

		pipeReader, pipeWriter := io.Pipe()

		go fixRepositoriesFile(name, imageReader, pipeWriter)

		opts := docker.LoadImageOptions{
			InputStream: pipeReader,
		}
		if err := md.client.LoadImage(opts); err != nil {
			return err
		}

		image, err = md.client.InspectImage(name)
		if err != nil {
			return err
		}
	}

	response.Images = []*rpc.Image{
		&rpc.Image{
			Id:   request.Id,
			Type: "container",
			Size: uint64(image.Size) / 1024 / 1024,
		},
	}
	return nil
}

// fixRepositoriesFile changes the repo name to the mistify-image-service's
// assigned image id and tag to "latest" before it is loaded into docker
func fixRepositoriesFile(newName string, in io.Reader, out io.WriteCloser) {
	defer out.Close()
	tarReader := tar.NewReader(in)
	tarWriter := tar.NewWriter(out)
	defer tarWriter.Close()

	for {
		header, err := tarReader.Next()
		if err != nil {
			if err == io.EOF {
				return
			}
			log.WithField("error", err).Error("failed to get next tar header")
			return
		}

		switch header.Typeflag {
		case tar.TypeReg:
			// Update the image name and tag in the repositories file
			if header.Name == "repositories" {
				// Read the file and parse the JSON
				// {"reponame":{"tag":"hash"}}
				repoMap := map[string]map[string]string{}
				jsonDecoder := json.NewDecoder(tarReader)
				if err := jsonDecoder.Decode(&repoMap); err != nil {
					log.WithField("error", err).Error("failed to parse repositories json")
					return
				}

				// Should only be one key. Replace it with the new repo name
				if len(repoMap) != 1 {
					log.WithFields(log.Fields{
						"error":   errors.New("incorrect number of repos"),
						"repoMap": repoMap,
					}).Error("must be only one repo specified")
					return
				}
				for oldName := range repoMap {
					tagMap := repoMap[oldName]
					// Should only be one tag. Replace it with the new repo name
					if len(tagMap) != 1 {
						log.WithFields(log.Fields{
							"error":   errors.New("incorrect number of tags"),
							"repoMap": repoMap,
						}).Error("must be only one tag specified")
						return
					}
					for oldTag := range tagMap {
						// Only rename if the tag is not already "latest"
						if oldTag == "latest" {
							break
						}
						tagMap["latest"] = tagMap[oldTag]
						delete(tagMap, oldTag)
					}
					repoMap[newName] = tagMap
					delete(repoMap, oldName)
				}

				// Update the header
				outBytes, err := json.Marshal(repoMap)
				if err != nil {
					log.WithField("error", err).Error("failed to marshal repositories json")
					return
				}
				header.Size = int64(len(outBytes))
				header.ModTime = time.Now()

				// Write the new header and data
				if err := tarWriter.WriteHeader(header); err != nil {
					log.WithField("error", err).Error("failed to write repositories header")
					return
				}
				if _, err := tarWriter.Write(outBytes); err != nil {
					log.WithField("error", err).Error("failed to write repositories json")
					return
				}
				continue
			}
			fallthrough
		default:
			// Direct copy
			if err := tarWriter.WriteHeader(header); err != nil {
				log.WithField("error", err).Error("failed to write tar header")
				return
			}
			if _, err := io.Copy(tarWriter, tarReader); err != nil {
				log.WithField("error", err).Error("failed to copy tar body")
				return
			}
		}
	}
}

// DeleteImage deletes a Docker image
func (md *MDocker) DeleteImage(h *http.Request, request *rpc.ImageRequest, response *rpc.ImageResponse) error {
	image, err := md.client.InspectImage(request.Id)
	if err != nil {
		return err
	}

	opts := docker.RemoveImageOptions{}
	if err := md.client.RemoveImageExtended(request.Id, opts); err != nil {
		return err
	}

	response.Images = []*rpc.Image{
		&rpc.Image{
			Id:   request.Id,
			Type: "container",
			Size: uint64(image.Size) / 1024 / 1024,
		},
	}
	return nil
}
<commit_msg>Move where the old repo key is deleted<commit_after>package mdocker

import (
	"archive/tar"
	"bufio"
	"compress/gzip"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/fsouza/go-dockerclient"
	"github.com/mistifyio/mistify-agent/rpc"
)

// ListImages retrieves a list of Docker images
func (md *MDocker) ListImages(h *http.Request, request *rpc.ImageRequest, response *rpc.ImageResponse) error {
	opts := docker.ListImagesOptions{}

	apiImages, err := md.client.ListImages(opts)
	if err != nil {
		return err
	}
	images := make([]*rpc.Image, 0, len(apiImages))
	for _, ai := range apiImages {
		id, _ := docker.ParseRepositoryTag(ai.RepoTags[0])
		images = append(images, &rpc.Image{
			Id:   id,
			Type: "container",
			Size: uint64(ai.Size) / 1024 / 1024,
		})
	}

	response.Images = images
	return nil
}

// GetImage retrieves information about a specific Docker image
func (md *MDocker) GetImage(h *http.Request, request *rpc.ImageRequest, response *rpc.ImageResponse) error {
	image, err := md.client.InspectImage(request.Id)
	if err != nil {
		return err
	}

	response.Images = []*rpc.Image{
		&rpc.Image{
			Id:   request.Id,
			Type: "container",
			Size: uint64(image.Size) / 1024 / 1024,
		},
	}
	return nil
}

// LoadImage downloads a new container image from the image service and
// imports it into Docker
func (md *MDocker) LoadImage(h *http.Request, request *rpc.ImageRequest, response *rpc.ImageResponse) error {
	name := request.Id

	// Check if we already have the image to avoid unnecessary pulling
	image, err := md.client.InspectImage(name)
	if err != nil && err != docker.ErrNoSuchImage {
		return err
	}

	if image == nil {
		// Docker Import lets the image get renamed, but it strips metadata
		// (which includes any CMD that had been set). Docker Load doesn't let
		// the image get renamed and doesn't return the image id or name after
		// loading, but does preserve metadata. Since a rename to use the
		// image-service assigned id and metadata with CMD are both necessary,
		// the only way forward is to update the repositories file inside and
		// then load it.
		source := fmt.Sprintf("http://%s/images/%s/download", md.imageService, request.Id)
		resp, err := http.Get(source)
		if err != nil {
			return err
		}
		defer resp.Body.Close()

		if resp.StatusCode != http.StatusOK {
			return ErrorHTTPCode{
				Expected: http.StatusOK,
				Code:     resp.StatusCode,
				Source:   source,
			}
		}

		// Use a response buffer so the first few bytes can be peeked at for
		// file type detection. Uncompress the image if it is gzipped
		responseBuffer := bufio.NewReader(resp.Body)
		var imageReader io.Reader = responseBuffer
		filetypeBytes, err := responseBuffer.Peek(512)
		if err != nil {
			return err
		}
		if http.DetectContentType(filetypeBytes) == "application/x-gzip" {
			gzipReader, err := gzip.NewReader(responseBuffer)
			if err != nil {
				return err
			}
			defer gzipReader.Close()
			imageReader = gzipReader
		}

		pipeReader, pipeWriter := io.Pipe()

		go fixRepositoriesFile(name, imageReader, pipeWriter)

		opts := docker.LoadImageOptions{
			InputStream: pipeReader,
		}
		if err := md.client.LoadImage(opts); err != nil {
			return err
		}

		image, err = md.client.InspectImage(name)
		if err != nil {
			return err
		}
	}

	response.Images = []*rpc.Image{
		&rpc.Image{
			Id:   request.Id,
			Type: "container",
			Size: uint64(image.Size) / 1024 / 1024,
		},
	}
	return nil
}

// fixRepositoriesFile changes the repo name to the mistify-image-service's
// assigned image id and tag to "latest" before it is loaded into docker
func fixRepositoriesFile(newName string, in io.Reader, out io.WriteCloser) {
	defer out.Close()
	tarReader := tar.NewReader(in)
	tarWriter := tar.NewWriter(out)
	defer tarWriter.Close()

	for {
		header, err := tarReader.Next()
		if err != nil {
			if err == io.EOF {
				return
			}
			log.WithField("error", err).Error("failed to get next tar header")
			return
		}

		switch header.Typeflag {
		case tar.TypeReg:
			// Update the image name and tag in the repositories file
			if header.Name == "repositories" {
				// Read the file and parse the JSON
				// {"reponame":{"tag":"hash"}}
				repoMap := map[string]map[string]string{}
				jsonDecoder := json.NewDecoder(tarReader)
				if err := jsonDecoder.Decode(&repoMap); err != nil {
					log.WithField("error", err).Error("failed to parse repositories json")
					return
				}

				// Should only be one key. Replace it with the new repo name
				if len(repoMap) != 1 {
					log.WithFields(log.Fields{
						"error":   errors.New("incorrect number of repos"),
						"repoMap": repoMap,
					}).Error("must be only one repo specified")
					return
				}
				for oldName := range repoMap {
					tagMap := repoMap[oldName]
					delete(repoMap, oldName)

					// Should only be one tag. Replace it with the new repo name
					if len(tagMap) != 1 {
						log.WithFields(log.Fields{
							"error":   errors.New("incorrect number of tags"),
							"repoMap": repoMap,
						}).Error("must be only one tag specified")
						return
					}
					for oldTag := range tagMap {
						// Only rename if the tag is not already "latest"
						if oldTag == "latest" {
							break
						}
						tagMap["latest"] = tagMap[oldTag]
						delete(tagMap, oldTag)
					}
					repoMap[newName] = tagMap
				}

				// Update the header
				outBytes, err := json.Marshal(repoMap)
				if err != nil {
					log.WithField("error", err).Error("failed to marshal repositories json")
					return
				}
				header.Size = int64(len(outBytes))
				header.ModTime = time.Now()

				// Write the new header and data
				if err := tarWriter.WriteHeader(header); err != nil {
					log.WithField("error", err).Error("failed to write repositories header")
					return
				}
				if _, err := tarWriter.Write(outBytes); err != nil {
					log.WithField("error", err).Error("failed to write repositories json")
					return
				}
				continue
			}
			fallthrough
		default:
			// Direct copy
			if err := tarWriter.WriteHeader(header); err != nil {
				log.WithField("error", err).Error("failed to write tar header")
				return
			}
			if _, err := io.Copy(tarWriter, tarReader); err != nil {
				log.WithField("error", err).Error("failed to copy tar body")
				return
			}
		}
	}
}

// DeleteImage deletes a Docker image
func (md *MDocker) DeleteImage(h *http.Request, request *rpc.ImageRequest, response *rpc.ImageResponse) error {
	image, err := md.client.InspectImage(request.Id)
	if err != nil {
		return err
	}

	opts := docker.RemoveImageOptions{}
	if err := md.client.RemoveImageExtended(request.Id, opts); err != nil {
		return err
	}

	response.Images = []*rpc.Image{
		&rpc.Image{
			Id:   request.Id,
			Type: "container",
			Size: uint64(image.Size) / 1024 / 1024,
		},
	}
	return nil
}
<|endoftext|>