<commit_before>// Copyright 2017 fatedier, fatedier@gmail.com
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package client

import (
	"context"
	"crypto/tls"
	"fmt"
	"io"
	"net"
	"runtime/debug"
	"sync"
	"time"

	"github.com/fatedier/frp/client/proxy"
	"github.com/fatedier/frp/pkg/auth"
	"github.com/fatedier/frp/pkg/config"
	"github.com/fatedier/frp/pkg/msg"
	"github.com/fatedier/frp/pkg/transport"
	frpNet "github.com/fatedier/frp/pkg/util/net"
	"github.com/fatedier/frp/pkg/util/xlog"

	"github.com/fatedier/golib/control/shutdown"
	"github.com/fatedier/golib/crypto"
	fmux "github.com/hashicorp/yamux"
)

type Control struct {
	// uniq id got from frps, attach it in loginMsg
	runID string

	// manage all proxies
	pxyCfgs map[string]config.ProxyConf
	pm      *proxy.Manager

	// manage all visitors
	vm *VisitorManager

	// control connection
	conn net.Conn

	// tcp stream multiplexing, if enabled
	session *fmux.Session

	// put a message in this channel to send it over control connection to server
	sendCh chan (msg.Message)

	// read from this channel to get the next message sent by server
	readCh chan (msg.Message)

	// goroutines can block by reading from this channel, it will be closed only in reader() when control connection is closed
	closedCh chan struct{}

	closedDoneCh chan struct{}

	// last time got the Pong message
	lastPong time.Time

	// The client configuration
	clientCfg config.ClientCommonConf

	readerShutdown     *shutdown.Shutdown
	writerShutdown     *shutdown.Shutdown
	msgHandlerShutdown *shutdown.Shutdown

	// The UDP port that the server is listening on
	serverUDPPort int

	mu sync.RWMutex

	xl *xlog.Logger

	// service context
	ctx context.Context

	// sets authentication based on selected method
	authSetter auth.Setter
}

func NewControl(ctx context.Context, runID string, conn net.Conn, session *fmux.Session,
	clientCfg config.ClientCommonConf,
	pxyCfgs map[string]config.ProxyConf,
	visitorCfgs map[string]config.VisitorConf,
	serverUDPPort int,
	authSetter auth.Setter) *Control {

	// new xlog instance
	ctl := &Control{
		runID:              runID,
		conn:               conn,
		session:            session,
		pxyCfgs:            pxyCfgs,
		sendCh:             make(chan msg.Message, 100),
		readCh:             make(chan msg.Message, 100),
		closedCh:           make(chan struct{}),
		closedDoneCh:       make(chan struct{}),
		clientCfg:          clientCfg,
		readerShutdown:     shutdown.New(),
		writerShutdown:     shutdown.New(),
		msgHandlerShutdown: shutdown.New(),
		serverUDPPort:      serverUDPPort,
		xl:                 xlog.FromContextSafe(ctx),
		ctx:                ctx,
		authSetter:         authSetter,
	}
	ctl.pm = proxy.NewManager(ctl.ctx, ctl.sendCh, clientCfg, serverUDPPort)

	ctl.vm = NewVisitorManager(ctl.ctx, ctl)
	ctl.vm.Reload(visitorCfgs)
	return ctl
}

func (ctl *Control) Run() {
	go ctl.worker()

	// start all proxies
	ctl.pm.Reload(ctl.pxyCfgs)

	// start all visitors
	go ctl.vm.Run()
	return
}

func (ctl *Control) HandleReqWorkConn(inMsg *msg.ReqWorkConn) {
	xl := ctl.xl
	workConn, err := ctl.connectServer()
	if err != nil {
		return
	}

	m := &msg.NewWorkConn{
		RunID: ctl.runID,
	}
	if err = ctl.authSetter.SetNewWorkConn(m); err != nil {
		xl.Warn("error during NewWorkConn authentication: %v", err)
		return
	}
	if err = msg.WriteMsg(workConn, m); err != nil {
		xl.Warn("work connection write to server error: %v", err)
		workConn.Close()
		return
	}

	var startMsg msg.StartWorkConn
	if err = msg.ReadMsgInto(workConn, &startMsg); err != nil {
		xl.Error("work connection closed before response StartWorkConn message: %v", err)
		workConn.Close()
		return
	}
	if startMsg.Error != "" {
		xl.Error("StartWorkConn contains error: %s", startMsg.Error)
		workConn.Close()
		return
	}

	// dispatch this work connection to related proxy
	ctl.pm.HandleWorkConn(startMsg.ProxyName, workConn, &startMsg)
}

func (ctl *Control) HandleNewProxyResp(inMsg *msg.NewProxyResp) {
	xl := ctl.xl
	// Server will return NewProxyResp message to each NewProxy message.
	// Start a new proxy handler if no error got
	err := ctl.pm.StartProxy(inMsg.ProxyName, inMsg.RemoteAddr, inMsg.Error)
	if err != nil {
		xl.Warn("[%s] start error: %v", inMsg.ProxyName, err)
	} else {
		xl.Info("[%s] start proxy success", inMsg.ProxyName)
	}
}

func (ctl *Control) Close() error {
	ctl.pm.Close()
	ctl.conn.Close()
	ctl.vm.Close()
	if ctl.session != nil {
		ctl.session.Close()
	}
	return nil
}

// ClosedDoneCh returns a channel which will be closed after all resources are released
func (ctl *Control) ClosedDoneCh() <-chan struct{} {
	return ctl.closedDoneCh
}

// connectServer return a new connection to frps
func (ctl *Control) connectServer() (conn net.Conn, err error) {
	xl := ctl.xl
	if ctl.clientCfg.TCPMux {
		stream, errRet := ctl.session.OpenStream()
		if errRet != nil {
			err = errRet
			xl.Warn("start new connection to server error: %v", err)
			return
		}
		conn = stream
	} else {
		var tlsConfig *tls.Config

		if ctl.clientCfg.TLSEnable {
			tlsConfig, err = transport.NewServerTLSConfig(
				ctl.clientCfg.TLSCertFile,
				ctl.clientCfg.TLSKeyFile,
				ctl.clientCfg.TLSTrustedCaFile)

			if err != nil {
				xl.Warn("fail to build tls configuration when connecting to server, err: %v", err)
				return
			}
		}
		conn, err = frpNet.ConnectServerByProxyWithTLS(ctl.clientCfg.HTTPProxy, ctl.clientCfg.Protocol,
			fmt.Sprintf("%s:%d", ctl.clientCfg.ServerAddr, ctl.clientCfg.ServerPort), tlsConfig)
		if err != nil {
			xl.Warn("start new connection to server error: %v", err)
			return
		}
	}
	return
}

// reader read all messages from frps and send to readCh
func (ctl *Control) reader() {
	xl := ctl.xl
	defer func() {
		if err := recover(); err != nil {
			xl.Error("panic error: %v", err)
			xl.Error(string(debug.Stack()))
		}
	}()
	defer ctl.readerShutdown.Done()
	defer close(ctl.closedCh)

	encReader := crypto.NewReader(ctl.conn, []byte(ctl.clientCfg.Token))
	for {
		m, err := msg.ReadMsg(encReader)
		if err != nil {
			if err == io.EOF {
				xl.Debug("read from control connection EOF")
				return
			}
			xl.Warn("read error: %v", err)
			ctl.conn.Close()
			return
		}
		ctl.readCh <- m
	}
}

// writer writes messages got from sendCh to frps
func (ctl *Control) writer() {
	xl := ctl.xl
	defer ctl.writerShutdown.Done()
	encWriter, err := crypto.NewWriter(ctl.conn, []byte(ctl.clientCfg.Token))
	if err != nil {
		xl.Error("crypto new writer error: %v", err)
		ctl.conn.Close()
		return
	}
	for {
		m, ok := <-ctl.sendCh
		if !ok {
			xl.Info("control writer is closing")
			return
		}

		if err := msg.WriteMsg(encWriter, m); err != nil {
			xl.Warn("write message to control connection error: %v", err)
			return
		}
	}
}

// msgHandler handles all channel events and do corresponding operations.
func (ctl *Control) msgHandler() {
	xl := ctl.xl
	defer func() {
		if err := recover(); err != nil {
			xl.Error("panic error: %v", err)
			xl.Error(string(debug.Stack()))
		}
	}()
	defer ctl.msgHandlerShutdown.Done()

	hbSend := time.NewTicker(time.Duration(ctl.clientCfg.HeartBeatInterval) * time.Second)
	defer hbSend.Stop()
	hbCheck := time.NewTicker(time.Second)
	defer hbCheck.Stop()

	ctl.lastPong = time.Now()

	for {
		select {
		case <-hbSend.C:
			// send heartbeat to server
			xl.Debug("send heartbeat to server")
			pingMsg := &msg.Ping{}
			if err := ctl.authSetter.SetPing(pingMsg); err != nil {
				xl.Warn("error during ping authentication: %v", err)
				return
			}
			ctl.sendCh <- pingMsg
		case <-hbCheck.C:
			if time.Since(ctl.lastPong) > time.Duration(ctl.clientCfg.HeartBeatTimeout)*time.Second {
				xl.Warn("heartbeat timeout")
				// let reader() stop
				ctl.conn.Close()
				return
			}
		case rawMsg, ok := <-ctl.readCh:
			if !ok {
				return
			}

			switch m := rawMsg.(type) {
			case *msg.ReqWorkConn:
				go ctl.HandleReqWorkConn(m)
			case *msg.NewProxyResp:
				ctl.HandleNewProxyResp(m)
			case *msg.Pong:
				if m.Error != "" {
					xl.Error("Pong contains error: %s", m.Error)
					ctl.conn.Close()
					return
				}
				ctl.lastPong = time.Now()
				xl.Debug("receive heartbeat from server")
			}
		}
	}
}

// If controler is notified by closedCh, reader and writer and handler will exit
func (ctl *Control) worker() {
	go ctl.msgHandler()
	go ctl.reader()
	go ctl.writer()

	select {
	case <-ctl.closedCh:
		// close related channels and wait until other goroutines done
		close(ctl.readCh)
		ctl.readerShutdown.WaitDone()
		ctl.msgHandlerShutdown.WaitDone()

		close(ctl.sendCh)
		ctl.writerShutdown.WaitDone()

		ctl.pm.Close()
		ctl.vm.Close()

		close(ctl.closedDoneCh)
		if ctl.session != nil {
			ctl.session.Close()
		}
		return
	}
}

func (ctl *Control) ReloadConf(pxyCfgs map[string]config.ProxyConf, visitorCfgs map[string]config.VisitorConf) error {
	ctl.vm.Reload(visitorCfgs)
	ctl.pm.Reload(pxyCfgs)
	return nil
}
<commit_msg>fix create tls work connection (#2013)<commit_after>// Copyright 2017 fatedier, fatedier@gmail.com
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package client

import (
	"context"
	"crypto/tls"
	"fmt"
	"io"
	"net"
	"runtime/debug"
	"sync"
	"time"

	"github.com/fatedier/frp/client/proxy"
	"github.com/fatedier/frp/pkg/auth"
	"github.com/fatedier/frp/pkg/config"
	"github.com/fatedier/frp/pkg/msg"
	"github.com/fatedier/frp/pkg/transport"
	frpNet "github.com/fatedier/frp/pkg/util/net"
	"github.com/fatedier/frp/pkg/util/xlog"

	"github.com/fatedier/golib/control/shutdown"
	"github.com/fatedier/golib/crypto"
	fmux "github.com/hashicorp/yamux"
)

type Control struct {
	// uniq id got from frps, attach it in loginMsg
	runID string

	// manage all proxies
	pxyCfgs map[string]config.ProxyConf
	pm      *proxy.Manager

	// manage all visitors
	vm *VisitorManager

	// control connection
	conn net.Conn

	// tcp stream multiplexing, if enabled
	session *fmux.Session

	// put a message in this channel to send it over control connection to server
	sendCh chan (msg.Message)

	// read from this channel to get the next message sent by server
	readCh chan (msg.Message)

	// goroutines can block by reading from this channel, it will be closed only in reader() when control connection is closed
	closedCh chan struct{}

	closedDoneCh chan struct{}

	// last time got the Pong message
	lastPong time.Time

	// The client configuration
	clientCfg config.ClientCommonConf

	readerShutdown     *shutdown.Shutdown
	writerShutdown     *shutdown.Shutdown
	msgHandlerShutdown *shutdown.Shutdown

	// The UDP port that the server is listening on
	serverUDPPort int

	mu sync.RWMutex

	xl *xlog.Logger

	// service context
	ctx context.Context

	// sets authentication based on selected method
	authSetter auth.Setter
}

func NewControl(ctx context.Context, runID string, conn net.Conn, session *fmux.Session,
	clientCfg config.ClientCommonConf,
	pxyCfgs map[string]config.ProxyConf,
	visitorCfgs map[string]config.VisitorConf,
	serverUDPPort int,
	authSetter auth.Setter) *Control {

	// new xlog instance
	ctl := &Control{
		runID:              runID,
		conn:               conn,
		session:            session,
		pxyCfgs:            pxyCfgs,
		sendCh:             make(chan msg.Message, 100),
		readCh:             make(chan msg.Message, 100),
		closedCh:           make(chan struct{}),
		closedDoneCh:       make(chan struct{}),
		clientCfg:          clientCfg,
		readerShutdown:     shutdown.New(),
		writerShutdown:     shutdown.New(),
		msgHandlerShutdown: shutdown.New(),
		serverUDPPort:      serverUDPPort,
		xl:                 xlog.FromContextSafe(ctx),
		ctx:                ctx,
		authSetter:         authSetter,
	}
	ctl.pm = proxy.NewManager(ctl.ctx, ctl.sendCh, clientCfg, serverUDPPort)

	ctl.vm = NewVisitorManager(ctl.ctx, ctl)
	ctl.vm.Reload(visitorCfgs)
	return ctl
}

func (ctl *Control) Run() {
	go ctl.worker()

	// start all proxies
	ctl.pm.Reload(ctl.pxyCfgs)

	// start all visitors
	go ctl.vm.Run()
	return
}

func (ctl *Control) HandleReqWorkConn(inMsg *msg.ReqWorkConn) {
	xl := ctl.xl
	workConn, err := ctl.connectServer()
	if err != nil {
		return
	}

	m := &msg.NewWorkConn{
		RunID: ctl.runID,
	}
	if err = ctl.authSetter.SetNewWorkConn(m); err != nil {
		xl.Warn("error during NewWorkConn authentication: %v", err)
		return
	}
	if err = msg.WriteMsg(workConn, m); err != nil {
		xl.Warn("work connection write to server error: %v", err)
		workConn.Close()
		return
	}

	var startMsg msg.StartWorkConn
	if err = msg.ReadMsgInto(workConn, &startMsg); err != nil {
		xl.Error("work connection closed before response StartWorkConn message: %v", err)
		workConn.Close()
		return
	}
	if startMsg.Error != "" {
		xl.Error("StartWorkConn contains error: %s", startMsg.Error)
		workConn.Close()
		return
	}

	// dispatch this work connection to related proxy
	ctl.pm.HandleWorkConn(startMsg.ProxyName, workConn, &startMsg)
}

func (ctl *Control) HandleNewProxyResp(inMsg *msg.NewProxyResp) {
	xl := ctl.xl
	// Server will return NewProxyResp message to each NewProxy message.
	// Start a new proxy handler if no error got
	err := ctl.pm.StartProxy(inMsg.ProxyName, inMsg.RemoteAddr, inMsg.Error)
	if err != nil {
		xl.Warn("[%s] start error: %v", inMsg.ProxyName, err)
	} else {
		xl.Info("[%s] start proxy success", inMsg.ProxyName)
	}
}

func (ctl *Control) Close() error {
	ctl.pm.Close()
	ctl.conn.Close()
	ctl.vm.Close()
	if ctl.session != nil {
		ctl.session.Close()
	}
	return nil
}

// ClosedDoneCh returns a channel which will be closed after all resources are released
func (ctl *Control) ClosedDoneCh() <-chan struct{} {
	return ctl.closedDoneCh
}

// connectServer return a new connection to frps
func (ctl *Control) connectServer() (conn net.Conn, err error) {
	xl := ctl.xl
	if ctl.clientCfg.TCPMux {
		stream, errRet := ctl.session.OpenStream()
		if errRet != nil {
			err = errRet
			xl.Warn("start new connection to server error: %v", err)
			return
		}
		conn = stream
	} else {
		var tlsConfig *tls.Config

		if ctl.clientCfg.TLSEnable {
			tlsConfig, err = transport.NewClientTLSConfig(
				ctl.clientCfg.TLSCertFile,
				ctl.clientCfg.TLSKeyFile,
				ctl.clientCfg.TLSTrustedCaFile,
				ctl.clientCfg.ServerAddr)

			if err != nil {
				xl.Warn("fail to build tls configuration when connecting to server, err: %v", err)
				return
			}
		}
		conn, err = frpNet.ConnectServerByProxyWithTLS(ctl.clientCfg.HTTPProxy, ctl.clientCfg.Protocol,
			fmt.Sprintf("%s:%d", ctl.clientCfg.ServerAddr, ctl.clientCfg.ServerPort), tlsConfig)
		if err != nil {
			xl.Warn("start new connection to server error: %v", err)
			return
		}
	}
	return
}

// reader read all messages from frps and send to readCh
func (ctl *Control) reader() {
	xl := ctl.xl
	defer func() {
		if err := recover(); err != nil {
			xl.Error("panic error: %v", err)
			xl.Error(string(debug.Stack()))
		}
	}()
	defer ctl.readerShutdown.Done()
	defer close(ctl.closedCh)

	encReader := crypto.NewReader(ctl.conn, []byte(ctl.clientCfg.Token))
	for {
		m, err := msg.ReadMsg(encReader)
		if err != nil {
			if err == io.EOF {
				xl.Debug("read from control connection EOF")
				return
			}
			xl.Warn("read error: %v", err)
			ctl.conn.Close()
			return
		}
		ctl.readCh <- m
	}
}

// writer writes messages got from sendCh to frps
func (ctl *Control) writer() {
	xl := ctl.xl
	defer ctl.writerShutdown.Done()
	encWriter, err := crypto.NewWriter(ctl.conn, []byte(ctl.clientCfg.Token))
	if err != nil {
		xl.Error("crypto new writer error: %v", err)
		ctl.conn.Close()
		return
	}
	for {
		m, ok := <-ctl.sendCh
		if !ok {
			xl.Info("control writer is closing")
			return
		}

		if err := msg.WriteMsg(encWriter, m); err != nil {
			xl.Warn("write message to control connection error: %v", err)
			return
		}
	}
}

// msgHandler handles all channel events and do corresponding operations.
func (ctl *Control) msgHandler() {
	xl := ctl.xl
	defer func() {
		if err := recover(); err != nil {
			xl.Error("panic error: %v", err)
			xl.Error(string(debug.Stack()))
		}
	}()
	defer ctl.msgHandlerShutdown.Done()

	hbSend := time.NewTicker(time.Duration(ctl.clientCfg.HeartBeatInterval) * time.Second)
	defer hbSend.Stop()
	hbCheck := time.NewTicker(time.Second)
	defer hbCheck.Stop()

	ctl.lastPong = time.Now()

	for {
		select {
		case <-hbSend.C:
			// send heartbeat to server
			xl.Debug("send heartbeat to server")
			pingMsg := &msg.Ping{}
			if err := ctl.authSetter.SetPing(pingMsg); err != nil {
				xl.Warn("error during ping authentication: %v", err)
				return
			}
			ctl.sendCh <- pingMsg
		case <-hbCheck.C:
			if time.Since(ctl.lastPong) > time.Duration(ctl.clientCfg.HeartBeatTimeout)*time.Second {
				xl.Warn("heartbeat timeout")
				// let reader() stop
				ctl.conn.Close()
				return
			}
		case rawMsg, ok := <-ctl.readCh:
			if !ok {
				return
			}

			switch m := rawMsg.(type) {
			case *msg.ReqWorkConn:
				go ctl.HandleReqWorkConn(m)
			case *msg.NewProxyResp:
				ctl.HandleNewProxyResp(m)
			case *msg.Pong:
				if m.Error != "" {
					xl.Error("Pong contains error: %s", m.Error)
					ctl.conn.Close()
					return
				}
				ctl.lastPong = time.Now()
				xl.Debug("receive heartbeat from server")
			}
		}
	}
}

// If controler is notified by closedCh, reader and writer and handler will exit
func (ctl *Control) worker() {
	go ctl.msgHandler()
	go ctl.reader()
	go ctl.writer()

	select {
	case <-ctl.closedCh:
		// close related channels and wait until other goroutines done
		close(ctl.readCh)
		ctl.readerShutdown.WaitDone()
		ctl.msgHandlerShutdown.WaitDone()

		close(ctl.sendCh)
		ctl.writerShutdown.WaitDone()

		ctl.pm.Close()
		ctl.vm.Close()

		close(ctl.closedDoneCh)
		if ctl.session != nil {
			ctl.session.Close()
		}
		return
	}
}

func (ctl *Control) ReloadConf(pxyCfgs map[string]config.ProxyConf, visitorCfgs map[string]config.VisitorConf) error {
	ctl.vm.Reload(visitorCfgs)
	ctl.pm.Reload(pxyCfgs)
	return nil
}
<|endoftext|>


<commit_before>package client

import (
	"fmt"
	"io"
	"log"
	"net"
	"strings"
	"time"
)

/*
BTConn contains the information necessary to maintain a P2P
connection to a Peer according to the BitTorrent protocol.
*/
type BTConn struct{}

type Handler interface {
	StartListening(chan BTConn, error)
}

/*
BTService is a wrapper around a TCPListener, along with
other state information.
*/
type BTService struct {
	Listener  *net.TCPListener
	Listening bool
	CloseCh   chan bool
	Port      int
}

/*
NewBTService returns a closed BTService on a specified port.
*/
func NewBTService(port int) *BTService {
	s := &BTService{
		Listening: false,
		CloseCh:   make(chan bool, 1),
		Port:      port,
	}
	return s
}

/*
StartListening starts a TCP listening service on a goroutine.
*/
func (s *BTService) StartListening() (err error) {
	log.Println("Start listening")
	addr, _ := net.ResolveTCPAddr("tcp", fmt.Sprintf(":%d", s.Port))
	l, err := net.ListenTCP("tcp", addr)
	if err != nil {
		return err
	}
	s.Listener = l
	s.Listening = true

	go func() {
		for {
			select {
			case <-s.CloseCh:
				log.Println("Closing BitTorrent Service")
				s.Listener.Close()
				s.Listening = false
				return
			default:
			}

			l.SetDeadline(time.Now().Add(time.Nanosecond))
			conn, err := l.AcceptTCP()
			if err != nil {
				if !strings.Contains(err.Error(), "i/o timeout") {
					log.Println(err)
				}
				continue
			} else {
				log.Println(conn)
				go handleConnection(conn)
			}

		}
	}()

	return nil
}

type Handshake struct {
	Length            byte
	Name              string
	ReservedExtension []byte
	Hash              []byte
	PeerID            []byte
}

func (h *Handshake) String() string {
	return fmt.Sprintf("pstrlen: %d, name: %s, reserved extension: %x , hash: %x , peer id: %s", h.Length, h.Name, h.ReservedExtension, h.Hash, h.PeerID)
}

func handleConnection(c net.Conn) {
	log.Println("Handle Connection")
	defer c.Close()

	// First connection, assume handshake messsage
	// Get the protocol name length
	buf := make([]byte, 1)
	log.Println("Waiting to readfull")
	_, err := io.ReadFull(c, buf)
	if err != nil {
		log.Println("[HandleConnection] Error: ", err)
		return
	}
	pstrLen := int(buf[0])

	// Get the rest of the handshake message
	buf = make([]byte, pstrLen+48)
	_, err = io.ReadFull(c, buf)
	if err != nil {
		log.Println("[HandleConnection] Error: ", err)
		return
	}

	// Parse fields out of the message
	handshake := &Handshake{
		Length:            byte(pstrLen),
		Name:              string(buf[0:pstrLen]),
		ReservedExtension: buf[pstrLen : pstrLen+8],
		Hash:              buf[pstrLen+8 : pstrLen+8+20],
		PeerID:            buf[pstrLen+8+20 : pstrLen+8+20+20],
	}

	log.Printf("[HandleConnection] Handshake: %q", buf)
	log.Printf("%q", handshake)

	log.Printf("Writing byte\n")
	c.Write([]byte("pong"))

	return
}

/*
StopListening stops the TCP listener by sending to its Close channel.
*/
func (s *BTService) StopListening() (err error) {
	// TODO: Check that listener is actually on
	fmt.Println("StopListening")
	s.CloseCh <- true
	return nil
}
<commit_msg>Start to handle TCP messages through channels.<commit_after>package client

import (
	"fmt"
	"io"
	"log"
	"net"
	"strings"
	"time"
)

/*
BTConn contains the information necessary to maintain a P2P
connection to a Peer according to the BitTorrent protocol.
*/
type BTConn struct {
}

type Handler interface {
	StartListening(chan BTConn, error)
}

/*
BTService is a wrapper around a TCPListener, along with
other state information.
*/
type BTService struct {
	Listener  *net.TCPListener
	Listening bool
	CloseCh   chan bool
	Port      int
}

/*
NewBTService returns a closed BTService on a specified port.
*/
func NewBTService(port int) *BTService {
	s := &BTService{
		Listening: false,
		CloseCh:   make(chan bool, 1),
		Port:      port,
	}
	return s
}

/*
StartListening starts a TCP listening service on a goroutine.
*/
func (s *BTService) StartListening() (err error) {
	log.Println("Start listening")
	addr, _ := net.ResolveTCPAddr("tcp", fmt.Sprintf(":%d", s.Port))
	l, err := net.ListenTCP("tcp", addr)
	if err != nil {
		return err
	}
	s.Listener = l
	s.Listening = true

	hsChan := make(chan net.Conn, 1)
	msgChan := make(chan string, 1)

	go func() {
		go handleMessages(hsChan, msgChan)

		for {
			select {
			case <-s.CloseCh:
				log.Println("Closing BitTorrent Service")
				s.Listener.Close()
				s.Listening = false
				return
			default:
			}

			l.SetDeadline(time.Now().Add(time.Nanosecond))
			conn, err := l.AcceptTCP()
			if err != nil {
				if !strings.Contains(err.Error(), "i/o timeout") {
					log.Println(err)
				}
				continue
			} else {
				log.Println(conn)
				go handleConnection(conn, hsChan)
			}

		}
	}()

	return nil
}

type Handshake struct {
	Length            byte
	Name              string
	ReservedExtension []byte
	Hash              []byte
	PeerID            []byte
}

func (h *Handshake) String() string {
	return fmt.Sprintf("pstrlen: %d, name: %s, reserved extension: %x , hash: %x , peer id: %s", h.Length, h.Name, h.ReservedExtension, h.Hash, h.PeerID)
}

func handleMessages(hsChan <-chan net.Conn, msgChan <-chan string) {
	peers := make(map[net.Conn]string)

	for {
		select {
		case hs := <-hsChan:
			handleHandshake(hs)
			peers[hs] = "added"
			log.Printf("Writing byte\n")
			hs.Write([]byte("pong"))
			hs.Close()
		default:

		}
	}
}

func handleConnection(c net.Conn, hsChan chan<- net.Conn) {
	log.Println("Handle Connection")

	hsChan <- c

	//	handleHandshake(c)

	//	log.Printf("Writing byte\n")
	//	c.Write([]byte("pong"))

	return
}

func handleHandshake(c net.Conn) {
	// First connection, assume handshake messsage
	// Get the protocol name length
	buf := make([]byte, 1)
	log.Println("Waiting to readfull")
	_, err := io.ReadFull(c, buf)
	if err != nil {
		log.Println("[HandleConnection] Error: ", err)
		return
	}
	pstrLen := int(buf[0])

	// Get the rest of the handshake message
	buf = make([]byte, pstrLen+48)
	_, err = io.ReadFull(c, buf)
	if err != nil {
		log.Println("[HandleConnection] Error: ", err)
		return
	}

	// Parse fields out of the message
	handshake := &Handshake{
		Length:            byte(pstrLen),
		Name:              string(buf[0:pstrLen]),
		ReservedExtension: buf[pstrLen : pstrLen+8],
		Hash:              buf[pstrLen+8 : pstrLen+8+20],
		PeerID:            buf[pstrLen+8+20 : pstrLen+8+20+20],
	}

	log.Printf("[HandleConnection] Handshake: %q", buf)
	log.Printf("%q", handshake)

	return
}

/*
StopListening stops the TCP listener by sending to its Close channel.
*/
func (s *BTService) StopListening() (err error) {
	// TODO: Check that listener is actually on
	fmt.Println("StopListening")
	s.CloseCh <- true
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudformation"
)

// stack - holds all meaningful information about a particular stack.
type stack struct {
	name         string
	stackname    string
	template     string
	dependsOn    []interface{}
	dependents   []interface{}
	stackoutputs *cloudformation.DescribeStacksOutput
}

// State - struct for handling stack deploy/terminate statuses
var state = struct {
	pending  string
	failed   string
	complete string
}{
	complete: "complete",
	pending:  "pending",
	failed:   "failed",
}

// mutex - used to sync access to cross thread variables
var mutex = &sync.Mutex{}

// updateState - Locks cross channel object and updates value
func updateState(statusMap map[string]string, name string, status string) {
	Log(fmt.Sprintf("Updating Stack Status Map: %s - %s", name, status), level.debug)
	mutex.Lock()
	statusMap[name] = status
	mutex.Unlock()
}

// setStackName - sets the stackname with struct
func (s *stack) setStackName() {
	s.stackname = fmt.Sprintf("%s-%s", project, s.name)
}

func (s *stack) deploy(session *session.Session) error {

	err := s.deployTimeParser()
	if err != nil {
		return err
	}

	Log(fmt.Sprintf("Updated Template:\n%s", s.template), level.debug)

	svc := cloudformation.New(session)
	capability := "CAPABILITY_IAM"

	createParams := &cloudformation.CreateStackInput{
		StackName:       aws.String(s.stackname),
		DisableRollback: aws.Bool(true), // no rollback by default
		TemplateBody:    aws.String(s.template),
		Capabilities:    []*string{&capability},
	}

	Log(fmt.Sprintln("Calling [CreateStack] with parameters:", createParams), level.debug)
	if _, err := svc.CreateStack(createParams); err != nil {
		return errors.New(fmt.Sprintln("Deploying failed: ", err.Error()))

	}

	go verbose(s.stackname, "CREATE", session)
	describeStacksInput := &cloudformation.DescribeStacksInput{
		StackName: aws.String(s.stackname),
	}

	Log(fmt.Sprintln("Calling [WaitUntilStackCreateComplete] with parameters:", describeStacksInput), level.debug)
	if err := svc.WaitUntilStackCreateComplete(describeStacksInput); err != nil {
		// FIXME this works in so far that we wait until the stack is
		// completed and capture errors, but it doesn't really tail
		// cloudroamtion events.
		return err
	}

	Log(fmt.Sprintf("Deployment successful: [%s]", s.stackname), "info")

	return nil
}

func (s *stack) update(session *session.Session) error {
	svc := cloudformation.New(session)
	capability := "CAPABILITY_IAM"
	updateParams := &cloudformation.UpdateStackInput{
		StackName:    aws.String(s.stackname),
		TemplateBody: aws.String(s.template),
		Capabilities: []*string{&capability},
	}

	if s.stackExists(session) {
		Log("Stack exists, updating...", "info")

		Log(fmt.Sprintln("Calling [UpdateStack] with parameters:", updateParams), level.debug)
		_, err := svc.UpdateStack(updateParams)

		if err != nil {
			return errors.New(fmt.Sprintln("Update failed: ", err))
		}

		go verbose(s.stackname, "UPDATE", session)

		describeStacksInput := &cloudformation.DescribeStacksInput{
			StackName: aws.String(s.stackname),
		}
		Log(fmt.Sprintln("Calling [WaitUntilStackUpdateComplete] with parameters:", describeStacksInput), level.debug)
		if err := svc.WaitUntilStackUpdateComplete(describeStacksInput); err != nil {
			// FIXME this works in so far that we wait until the stack is
			// completed and capture errors, but it doesn't really tail
			// cloudroamtion events.
			return err
		}

		Log(fmt.Sprintf("Stack update successful: [%s]", s.stackname), "info")

	}
	return nil
}

func (s *stack) terminate(session *session.Session) error {
	svc := cloudformation.New(session)

	params := &cloudformation.DeleteStackInput{
		StackName: aws.String(s.stackname),
	}

	Log(fmt.Sprintln("Calling [DeleteStack] with parameters:", params), level.debug)
	_, err := svc.DeleteStack(params)

	go verbose(s.stackname, "DELETE", session)

	if err != nil {
		return errors.New(fmt.Sprintln("Deleting failed: ", err))
	}

	describeStacksInput := &cloudformation.DescribeStacksInput{
		StackName: aws.String(s.stackname),
	}

	Log(fmt.Sprintln("Calling [WaitUntilStackDeleteComplete] with parameters:", describeStacksInput), level.debug)
	if err := svc.WaitUntilStackDeleteComplete(describeStacksInput); err != nil {
		// FIXME this works in so far that we wait until the stack is
		// completed and capture errors, but it doesn't really tail
		// cloudroamtion events.
		return err
	}

	Log(fmt.Sprintf("Deletion successful: [%s]", s.stackname), "info")
	return nil
}

func (s *stack) stackExists(session *session.Session) bool {
	svc := cloudformation.New(session)

	describeStacksInput := &cloudformation.DescribeStacksInput{
		StackName: aws.String(s.stackname),
	}

	Log(fmt.Sprintln("Calling [DescribeStacks] with parameters:", describeStacksInput), level.debug)
	_, err := svc.DescribeStacks(describeStacksInput)

	if err == nil {
		return true
	}

	return false
}

func (s *stack) status(session *session.Session) error {
	svc := cloudformation.New(session)

	describeStacksInput := &cloudformation.DescribeStacksInput{
		StackName: aws.String(s.stackname),
	}

	Log(fmt.Sprintln("Calling [UpdateStack] with parameters:", describeStacksInput), level.debug)
	status, err := svc.DescribeStacks(describeStacksInput)

	if err != nil {
		if strings.Contains(strings.ToLower(err.Error()), "exist") {
			fmt.Printf("create_pending -> %s [%s]"+"\n", s.name, s.stackname)
			return nil
		}
		return err
	}

	// Define time flag
	stat := *status.Stacks[0].StackStatus
	var timeflag time.Time
	switch strings.Split(stat, "_")[0] {
	case "UPDATE":
		timeflag = *status.Stacks[0].LastUpdatedTime
	default:
		timeflag = *status.Stacks[0].CreationTime
	}

	// Print Status
	fmt.Printf(
		"%s%s - %s --> %s - [%s]"+"\n",
		colorString(`@`, "magenta"),
		timeflag.Format(time.RFC850),
		strings.ToLower(colorMap(*status.Stacks[0].StackStatus)),
		s.name,
		s.stackname,
	)

	return nil
}

// StackOutputs - Returns outputs of given stackname
func StackOutputs(name string, session *session.Session) (*cloudformation.DescribeStacksOutput, error) {

	svc := cloudformation.New(session)
	outputParams := &cloudformation.DescribeStacksInput{
		StackName: aws.String(name),
	}

	Log(fmt.Sprintln("Calling [DescribeStacks] with parameters:", outputParams), level.debug)
	outputs, err := svc.DescribeStacks(outputParams)
	if err != nil {
		return &cloudformation.DescribeStacksOutput{}, errors.New(fmt.Sprintln("Unable to reach stack", err.Error()))
	}

	return outputs, nil
}

// Exports - prints all cloudformation exports
func Exports(session *session.Session) error {

	svc := cloudformation.New(session)

	exportParams := &cloudformation.ListExportsInput{}

	Log(fmt.Sprintln("Calling [ListExports] with parameters:", exportParams), level.debug)
	exports, err := svc.ListExports(exportParams)

	if err != nil {
		return err
	}

	for _, i := range exports.Exports {

		fmt.Printf("Export Name: %s\nExport Value: %s\n--\n", colorString(*i.Name, "magenta"), *i.Value)
	}

	return nil
}

func (s *stack) state(session *session.Session) (string, error) {
	svc := cloudformation.New(session)

	describeStacksInput := &cloudformation.DescribeStacksInput{
		StackName: aws.String(s.stackname),
	}

	Log(fmt.Sprintln("Calling [DescribeStacks] with parameters: ", describeStacksInput), level.debug)
	status, err := svc.DescribeStacks(describeStacksInput)
	if err != nil {
		if strings.Contains(err.Error(), "not exist") {
			return state.pending, nil
		}
		return "", err
	}

	if strings.Contains(strings.ToLower(status.GoString()), "complete") {
		return state.complete, nil
	} else if strings.Contains(strings.ToLower(status.GoString()), "fail") {
		return state.failed, nil
	}
	return "", nil
}

// Check - Validates Cloudformation Templates
func Check(template string, session *session.Session) error {
	svc := cloudformation.New(session)
	params := &cloudformation.ValidateTemplateInput{
		TemplateBody: aws.String(template),
	}

	Log(fmt.Sprintf("Calling [ValidateTemplate] with parameters:\n%s"+"\n--\n", params), level.debug)
	resp, err := svc.ValidateTemplate(params)
	if err != nil {
		return err
	}

	fmt.Printf(
		"%s\n\n%s"+"\n",
		colorString("Valid!", "green"),
		resp.GoString(),
	)

	return nil

}

// DeployHandler - Handles deploying stacks in the corrcet order
func DeployHandler() {
	// status -  pending, failed, completed
	var status = make(map[string]string)

	sess, _ := awsSession()

	for _, stk := range stacks {
		// Set deploy status & Check if stack exists
		if stk.stackExists(sess) {

			updateState(status, stk.name, state.complete)
			fmt.Printf("Stack [%s] already exists..."+"\n", stk.name)
			continue
		} else {
			updateState(status, stk.name, state.pending)
		}

		if len(stk.dependsOn) == 0 {
			wg.Add(1)
			go func(s stack, sess *session.Session) {
				defer wg.Done()

				// Deploy 0 Depency Stacks first - each on their on go routine
				Log(fmt.Sprintf("Deploying a template for [%s]", s.name), "info")

				if err := s.deploy(sess); err != nil {
					handleError(err)
				}

				updateState(status, s.name, state.complete)

				// TODO: add deploy logic here
				return
			}(*stk, sess)
			continue
		}

		wg.Add(1)
		go func(s *stack, sess *session.Session) {
			Log(fmt.Sprintf("[%s] depends on: %s", s.name, s.dependsOn), "info")
			defer wg.Done()

			Log(fmt.Sprintf("Beginning Wait State for Depencies of [%s]"+"\n", s.name), level.debug)
			for {
				depts := []string{}
				for _, dept := range s.dependsOn {
					// Dependency wait
					dp := &stack{name: dept.(string)}
					dp.setStackName()
					chk, _ := dp.state(sess)

					switch chk {
					case state.failed:
						updateState(status, dp.name, state.failed)
					case state.complete:
						updateState(status, dp.name, state.complete)
					default:
						updateState(status, dp.name, state.pending)
					}

					mutex.Lock()
					depts = append(depts, status[dept.(string)])
					mutex.Unlock()
				}

				if all(depts, state.complete) {
					// Deploy stack once dependencies clear
					Log(fmt.Sprintf("Deploying a template for [%s]", s.name), "info")

					if err := s.deploy(sess); err != nil {
						handleError(err)
					}
					return
				}

				for _, v := range depts {
					if v == state.failed {
						Log(fmt.Sprintf("Deploy Cancelled for stack [%s] due to dependency failure!", s.name), "warn")
						return
					}
				}

				time.Sleep(time.Second * 1)
			}
		}(stk, sess)

	}

	// Wait for go routines to complete
	wg.Wait()
}

// TerminateHandler - Handles terminating stacks in the correct order
func TerminateHandler() {
	// status -  pending, failed, completed
	var status = make(map[string]string)

	sess, _ := awsSession()

	for _, stk := range stacks {
		// Check if stack exists

		if len(stk.dependsOn) == 0 {
			wg.Add(1)
			go func(s stack, sess *session.Session) {
				defer wg.Done()
				// Reverse depency look-up so termination waits for all stacks
				// which depend on it, to finish terminating first.
				for {
					depts := []string{}

					for _, stk := range stacks {

						if stringIn(s.name, stk.dependsOn) {

							mutex.Lock()
							depts = append(depts, status[stk.name])
							mutex.Unlock()
						}
					}

					if all(depts, state.complete) {
						s.terminate(sess)
						updateState(status, s.name, state.complete)
						return
					}

				}

			}(*stk, sess)
			continue
		}

		wg.Add(1)
		go func(s *stack, sess *session.Session) {
			defer wg.Done()

			// Stacks with no Reverse depencies are terminate first
			updateState(status, s.name, state.pending)

			Log(fmt.Sprintf("Terminating stack [%s]", s.stackname), "info")
			if err := s.terminate(sess); err != nil {
				updateState(status, s.name, state.failed)
				return
			}

			updateState(status, s.name, state.complete)

			return

		}(stk, sess)

	}

	// Wait for go routines to complete
	wg.Wait()
}
<commit_msg>fixed reverse depency termination bug<commit_after>package main

import (
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudformation"
)

// stack - holds all meaningful information about a particular stack.
type stack struct {
	name         string
	stackname    string
	template     string
	dependsOn    []interface{}
	dependents   []interface{}
	stackoutputs *cloudformation.DescribeStacksOutput
}

// State - struct for handling stack deploy/terminate statuses
var state = struct {
	pending  string
	failed   string
	complete string
}{
	complete: "complete",
	pending:  "pending",
	failed:   "failed",
}

// mutex - used to sync access to cross thread variables
var mutex = &sync.Mutex{}

// updateState - Locks cross channel object and updates value
func updateState(statusMap map[string]string, name string, status string) {
	Log(fmt.Sprintf("Updating Stack Status Map: %s - %s", name, status), level.debug)
	mutex.Lock()
	statusMap[name] = status
	mutex.Unlock()
}

// setStackName - sets the stackname with struct
func (s *stack) setStackName() {
	s.stackname = fmt.Sprintf("%s-%s", project, s.name)
}

func (s *stack) deploy(session *session.Session) error {

	err := s.deployTimeParser()
	if err != nil {
		return err
	}

	Log(fmt.Sprintf("Updated Template:\n%s", s.template), level.debug)

	svc := cloudformation.New(session)
	capability := "CAPABILITY_IAM"

	createParams := &cloudformation.CreateStackInput{
		StackName:       aws.String(s.stackname),
		DisableRollback: aws.Bool(true), // no rollback by default
		TemplateBody:    aws.String(s.template),
		Capabilities:    []*string{&capability},
	}

	Log(fmt.Sprintln("Calling [CreateStack] with parameters:", createParams), level.debug)
	if _, err := svc.CreateStack(createParams); err != nil {
		return errors.New(fmt.Sprintln("Deploying failed: ", err.Error()))

	}

	go verbose(s.stackname, "CREATE", session)
	describeStacksInput := &cloudformation.DescribeStacksInput{
		StackName: aws.String(s.stackname),
	}

	Log(fmt.Sprintln("Calling [WaitUntilStackCreateComplete] with parameters:", describeStacksInput), level.debug)
	if err := svc.WaitUntilStackCreateComplete(describeStacksInput); err != nil {
		// FIXME this works in so far that we wait until the stack is
		// completed and capture errors, but it doesn't really tail
		// cloudroamtion events.
		return err
	}

	Log(fmt.Sprintf("Deployment successful: [%s]", s.stackname), "info")

	return nil
}

func (s *stack) update(session *session.Session) error {
	svc := cloudformation.New(session)
	capability := "CAPABILITY_IAM"
	updateParams := &cloudformation.UpdateStackInput{
		StackName:    aws.String(s.stackname),
		TemplateBody: aws.String(s.template),
		Capabilities: []*string{&capability},
	}

	if s.stackExists(session) {
		Log("Stack exists, updating...", "info")

		Log(fmt.Sprintln("Calling [UpdateStack] with parameters:", updateParams), level.debug)
		_, err := svc.UpdateStack(updateParams)

		if err != nil {
			return errors.New(fmt.Sprintln("Update failed: ", err))
		}

		go verbose(s.stackname, "UPDATE", session)

		describeStacksInput := &cloudformation.DescribeStacksInput{
			StackName: aws.String(s.stackname),
		}
		Log(fmt.Sprintln("Calling [WaitUntilStackUpdateComplete] with parameters:", describeStacksInput), level.debug)
		if err := svc.WaitUntilStackUpdateComplete(describeStacksInput); err != nil {
			// FIXME this works in so far that we wait until the stack is
			// completed and capture errors, but it doesn't really tail
			// cloudroamtion events.
			return err
		}

		Log(fmt.Sprintf("Stack update successful: [%s]", s.stackname), "info")

	}
	return nil
}

func (s *stack) terminate(session *session.Session) error {
	svc := cloudformation.New(session)

	params := &cloudformation.DeleteStackInput{
		StackName: aws.String(s.stackname),
	}

	Log(fmt.Sprintln("Calling [DeleteStack] with parameters:", params), level.debug)
	_, err := svc.DeleteStack(params)

	go verbose(s.stackname, "DELETE", session)

	if err != nil {
		return errors.New(fmt.Sprintln("Deleting failed: ", err))
	}

	describeStacksInput := &cloudformation.DescribeStacksInput{
		StackName: aws.String(s.stackname),
	}

	Log(fmt.Sprintln("Calling [WaitUntilStackDeleteComplete] with parameters:", describeStacksInput), level.debug)
	if err := svc.WaitUntilStackDeleteComplete(describeStacksInput); err != nil {
		// FIXME this works in so far that we wait until the stack is
		// completed and capture errors, but it doesn't really tail
		// cloudroamtion events.
		return err
	}

	Log(fmt.Sprintf("Deletion successful: [%s]", s.stackname), "info")
	return nil
}

func (s *stack) stackExists(session *session.Session) bool {
	svc := cloudformation.New(session)

	describeStacksInput := &cloudformation.DescribeStacksInput{
		StackName: aws.String(s.stackname),
	}

	Log(fmt.Sprintln("Calling [DescribeStacks] with parameters:", describeStacksInput), level.debug)
	_, err := svc.DescribeStacks(describeStacksInput)

	if err == nil {
		return true
	}

	return false
}

func (s *stack) status(session *session.Session) error {
	svc := cloudformation.New(session)

	describeStacksInput := &cloudformation.DescribeStacksInput{
		StackName: aws.String(s.stackname),
	}

	Log(fmt.Sprintln("Calling [UpdateStack] with parameters:", describeStacksInput), level.debug)
	status, err := svc.DescribeStacks(describeStacksInput)

	if err != nil {
		if strings.Contains(strings.ToLower(err.Error()), "exist") {
			fmt.Printf("create_pending -> %s [%s]"+"\n", s.name, s.stackname)
			return nil
		}
		return err
	}

	// Define time flag
	stat := *status.Stacks[0].StackStatus
	var timeflag time.Time
	switch strings.Split(stat, "_")[0] {
	case "UPDATE":
		timeflag = *status.Stacks[0].LastUpdatedTime
	default:
		timeflag = *status.Stacks[0].CreationTime
	}

	// Print Status
	fmt.Printf(
		"%s%s - %s --> %s - [%s]"+"\n",
		colorString(`@`, "magenta"),
		timeflag.Format(time.RFC850),
		strings.ToLower(colorMap(*status.Stacks[0].StackStatus)),
		s.name,
		s.stackname,
	)

	return nil
}

// StackOutputs - Returns outputs of given stackname
func StackOutputs(name string, session *session.Session) (*cloudformation.DescribeStacksOutput, error) {

	svc := cloudformation.New(session)
	outputParams := &cloudformation.DescribeStacksInput{
		StackName: aws.String(name),
	}

	Log(fmt.Sprintln("Calling [DescribeStacks] with parameters:", outputParams), level.debug)
	outputs, err := svc.DescribeStacks(outputParams)
	if err != nil {
		return &cloudformation.DescribeStacksOutput{}, errors.New(fmt.Sprintln("Unable to reach stack", err.Error()))
	}

	return outputs, nil
}

// Exports - prints all cloudformation exports
func Exports(session *session.Session) error {

	svc := cloudformation.New(session)

	exportParams := &cloudformation.ListExportsInput{}

	Log(fmt.Sprintln("Calling [ListExports] with parameters:", exportParams), level.debug)
	exports, err := svc.ListExports(exportParams)

	if err != nil {
		return err
	}

	for _, i := range exports.Exports {

		fmt.Printf("Export Name: %s\nExport Value: %s\n--\n", colorString(*i.Name, "magenta"), *i.Value)
	}

	return nil
}

func (s *stack) state(session *session.Session) (string, error) {
	svc := cloudformation.New(session)

	describeStacksInput := &cloudformation.DescribeStacksInput{
		StackName: aws.String(s.stackname),
	}

	Log(fmt.Sprintln("Calling [DescribeStacks] with parameters: ", describeStacksInput), level.debug)
	status, err := svc.DescribeStacks(describeStacksInput)
	if err != nil {
		if strings.Contains(err.Error(), "not exist") {
			return state.pending, nil
		}
		return "", err
	}

	if strings.Contains(strings.ToLower(status.GoString()), "complete") {
		return state.complete, nil
	} else if strings.Contains(strings.ToLower(status.GoString()), "fail") {
		return state.failed, nil
	}
	return "", nil
}

// Check - Validates Cloudformation Templates
func Check(template string, session *session.Session) error {
	svc := cloudformation.New(session)
	params := &cloudformation.ValidateTemplateInput{
		TemplateBody: aws.String(template),
	}

	Log(fmt.Sprintf("Calling [ValidateTemplate] with parameters:\n%s"+"\n--\n", params), level.debug)
	resp, err := svc.ValidateTemplate(params)
	if err != nil {
		return err
	}

	fmt.Printf(
		"%s\n\n%s"+"\n",
		colorString("Valid!", "green"),
		resp.GoString(),
	)

	return nil

}

// DeployHandler - Handles deploying stacks in the corrcet order
func DeployHandler() {
	// status -  pending, failed, completed
	var status = make(map[string]string)

	sess, _ := awsSession()

	for _, stk := range stacks {

		if _, ok := job.stacks[stk.name]; !ok && len(job.stacks) > 0 {
			continue
		}

		// Set deploy status & Check if stack exists
		if stk.stackExists(sess) {

			updateState(status, stk.name, state.complete)
			fmt.Printf("Stack [%s] already exists..."+"\n", stk.name)
			continue
		} else {
			updateState(status, stk.name, state.pending)
		}

		if len(stk.dependsOn) == 0 {
			wg.Add(1)
			go func(s stack, sess *session.Session) {
				defer wg.Done()

				// Deploy 0 Depency Stacks first - each on their on go routine
				Log(fmt.Sprintf("Deploying a template for [%s]", s.name), "info")

				if err := s.deploy(sess); err != nil {
					handleError(err)
				}

				updateState(status, s.name, state.complete)

				// TODO: add deploy logic here
				return
			}(*stk, sess)
			continue
		}

		wg.Add(1)
		go func(s *stack, sess *session.Session) {
			Log(fmt.Sprintf("[%s] depends on: %s", s.name, s.dependsOn), "info")
			defer wg.Done()

			Log(fmt.Sprintf("Beginning Wait State for Depencies of [%s]"+"\n", s.name), level.debug)
			for {
				depts := []string{}
				for _, dept := range s.dependsOn {
					// Dependency wait
					dp := &stack{name: dept.(string)}
					dp.setStackName()
					chk, _ := dp.state(sess)

					switch chk {
					case state.failed:
						updateState(status, dp.name, state.failed)
					case state.complete:
						updateState(status, dp.name, state.complete)
					default:
						updateState(status, dp.name, state.pending)
					}

					mutex.Lock()
					depts = append(depts, status[dept.(string)])
					mutex.Unlock()
				}

				if all(depts, state.complete) {
					// Deploy stack once dependencies clear
					Log(fmt.Sprintf("Deploying a template for [%s]", s.name), "info")

					if err := s.deploy(sess); err != nil {
						handleError(err)
					}
					return
				}

				for _, v := range depts {
					if v == state.failed {
						Log(fmt.Sprintf("Deploy Cancelled for stack [%s] due to dependency failure!", s.name), "warn")
						return
					}
				}

				time.Sleep(time.Second * 1)
			}
		}(stk, sess)

	}

	// Wait for go routines to complete
	wg.Wait()
}

// TerminateHandler - Handles terminating stacks in the correct order
func TerminateHandler() {
	// 	status -  pending, failed, completed
	var status = make(map[string]string)

	sess, _ := awsSession()

	for _, stk := range stacks {
		if _, ok := job.stacks[stk.name]; !ok && len(job.stacks) > 0 {
			Log(fmt.Sprintf("%s: not in job.stacks, skipping", stk.name), level.debug)
			continue // only process items in the job.stacks unless empty
		}

		if len(stk.dependsOn) == 0 {
			wg.Add(1)
			go func(s stack, sess *session.Session) {
				defer wg.Done()
				// Reverse depency look-up so termination waits for all stacks
				// which depend on it, to finish terminating first.
				for {

					for _, stk := range stacks {
						// fmt.Println(stk, stk.dependsOn)
						if stringIn(s.name, stk.dependsOn) {
							Log(fmt.Sprintf("[%s]: Depends on [%s].. Waiting for dependency to terminate", stk.name, s.name), level.info)
							for {

								if !stk.stackExists(sess) {
									break
								}
								time.Sleep(time.Second * 2)
							}
						}
					}

					s.terminate(sess)
					return
				}

			}(*stk, sess)
			continue
		}

		wg.Add(1)
		go func(s *stack, sess *session.Session) {
			defer wg.Done()

			// Stacks with no Reverse depencies are terminated first
			updateState(status, s.name, state.pending)

			Log(fmt.Sprintf("Terminating stack [%s]", s.stackname), "info")
			if err := s.terminate(sess); err != nil {
				updateState(status, s.name, state.failed)
				return
			}

			updateState(status, s.name, state.complete)

			return

		}(stk, sess)

	}

	// Wait for go routines to complete
	wg.Wait()
}
<|endoftext|>


<commit_before>package cluster

import (
	"errors"
	"sync"

	"github.com/flynn/flynn-host/types"
	"github.com/flynn/go-discoverd"
	"github.com/flynn/rpcplus"
)

var ErrNoServers = errors.New("cluster: no servers found")

func NewClient() (*Client, error) {
	services, err := discoverd.NewServiceSet("flynn-host")
	client := &Client{service: services}
	go client.followLeader()
	return client, err
}

func (c *Client) followLeader() {
	for update := range c.service.Leaders() {
		c.mtx.Lock()
		if closer, ok := c.c.(interface {
			Close() error
		}); ok {
			closer.Close()
		}
		c.c, c.err = rpcplus.DialHTTP("tcp", update.Addr)
		// TODO: use attempt package to retry here
		c.mtx.Unlock()
	}
	// TODO: reconnect to discoverd here
}

type Client struct {
	service discoverd.ServiceSet

	c   RPCClient
	mtx sync.RWMutex
	err error
}

func (c *Client) ListHosts() (map[string]host.Host, error) {
	c.mtx.RLock()
	if err := c.err; err != nil {
		c.mtx.RUnlock()
		return nil, err
	}
	client := c.c
	c.mtx.RUnlock()

	var state map[string]host.Host
	err := client.Call("Cluster.ListHosts", struct{}{}, &state)
	return state, err
}

func (c *Client) AddJobs(req *host.AddJobsReq) (*host.AddJobsRes, error) {
	c.mtx.RLock()
	if err := c.err; err != nil {
		c.mtx.RUnlock()
		return nil, err
	}
	client := c.c
	c.mtx.RUnlock()

	var res host.AddJobsRes
	err := client.Call("Cluster.AddJobs", req, &res)
	return &res, err
}

func (c *Client) ConnectHost(id string) (*Host, error) {
	// TODO: reuse connection if leader id == id
	services := c.service.Select(map[string]string{"id": id})
	if len(services) == 0 {
		return nil, ErrNoServers
	}
	rc, err := rpcplus.DialHTTP("tcp", services[0].Addr)
	return &Host{service: c.service, c: rc}, err
}

func (c *Client) RPCClient() (RPCClient, error) {
	c.mtx.RLock()
	defer c.mtx.RUnlock()
	return c.c, c.err
}

type Host struct {
	service discoverd.ServiceSet

	c RPCClient
}

type RPCClient interface {
	Call(serviceMethod string, args interface{}, reply interface{}) error
	Go(serviceMethod string, args interface{}, reply interface{}, done chan *rpcplus.Call) *rpcplus.Call
	StreamGo(serviceMethod string, args interface{}, replyStream interface{}) *rpcplus.Call
}

func (c *Host) ListJobs() (map[string]host.ActiveJob, error) {
	var jobs map[string]host.ActiveJob
	err := c.c.Call("Host.ListJobs", struct{}{}, &jobs)
	return jobs, err
}

func (c *Host) GetJob(id string) (*host.ActiveJob, error) {
	var res host.ActiveJob
	err := c.c.Call("Host.GetJob", id, &res)
	return &res, err
}

func (c *Host) StopJob(id string) error {
	return c.c.Call("Host.StopJob", id, &struct{}{})
}
<commit_msg>pkg/cluster: Don't return from NewClient until connected<commit_after>package cluster

import (
	"errors"
	"sync"

	"github.com/flynn/flynn-host/types"
	"github.com/flynn/go-discoverd"
	"github.com/flynn/rpcplus"
)

var ErrNoServers = errors.New("cluster: no servers found")

func NewClient() (*Client, error) {
	services, err := discoverd.NewServiceSet("flynn-host")
	if err != nil {
		return nil, err
	}
	client := &Client{service: services}
	firstErr := make(chan error)
	go client.followLeader(firstErr)
	return client, <-firstErr
}

func (c *Client) followLeader(firstErr chan<- error) {
	for update := range c.service.Leaders() {
		c.mtx.Lock()
		if closer, ok := c.c.(interface {
			Close() error
		}); ok {
			closer.Close()
		}
		c.c, c.err = rpcplus.DialHTTP("tcp", update.Addr)
		// TODO: use attempt package to retry here
		c.mtx.Unlock()
		if firstErr != nil {
			firstErr <- c.err
			firstErr = nil
		}
	}
	// TODO: reconnect to discoverd here
}

type Client struct {
	service discoverd.ServiceSet

	c   RPCClient
	mtx sync.RWMutex
	err error
}

func (c *Client) ListHosts() (map[string]host.Host, error) {
	c.mtx.RLock()
	if err := c.err; err != nil {
		c.mtx.RUnlock()
		return nil, err
	}
	client := c.c
	c.mtx.RUnlock()

	var state map[string]host.Host
	err := client.Call("Cluster.ListHosts", struct{}{}, &state)
	return state, err
}

func (c *Client) AddJobs(req *host.AddJobsReq) (*host.AddJobsRes, error) {
	c.mtx.RLock()
	if err := c.err; err != nil {
		c.mtx.RUnlock()
		return nil, err
	}
	client := c.c
	c.mtx.RUnlock()

	var res host.AddJobsRes
	err := client.Call("Cluster.AddJobs", req, &res)
	return &res, err
}

func (c *Client) ConnectHost(id string) (*Host, error) {
	// TODO: reuse connection if leader id == id
	services := c.service.Select(map[string]string{"id": id})
	if len(services) == 0 {
		return nil, ErrNoServers
	}
	rc, err := rpcplus.DialHTTP("tcp", services[0].Addr)
	return &Host{service: c.service, c: rc}, err
}

func (c *Client) RPCClient() (RPCClient, error) {
	c.mtx.RLock()
	defer c.mtx.RUnlock()
	return c.c, c.err
}

type Host struct {
	service discoverd.ServiceSet

	c RPCClient
}

type RPCClient interface {
	Call(serviceMethod string, args interface{}, reply interface{}) error
	Go(serviceMethod string, args interface{}, reply interface{}, done chan *rpcplus.Call) *rpcplus.Call
	StreamGo(serviceMethod string, args interface{}, replyStream interface{}) *rpcplus.Call
}

func (c *Host) ListJobs() (map[string]host.ActiveJob, error) {
	var jobs map[string]host.ActiveJob
	err := c.c.Call("Host.ListJobs", struct{}{}, &jobs)
	return jobs, err
}

func (c *Host) GetJob(id string) (*host.ActiveJob, error) {
	var res host.ActiveJob
	err := c.c.Call("Host.GetJob", id, &res)
	return &res, err
}

func (c *Host) StopJob(id string) error {
	return c.c.Call("Host.StopJob", id, &struct{}{})
}
<|endoftext|>


<commit_before>// Copyright © 2017 Alexander Sosna <alexander@xxor.de>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cmd

import (
	"io"
	"os"
	"os/exec"
	"sync"

	log "github.com/Sirupsen/logrus"
	ec "github.com/xxorde/pgglaskugel/errorcheck"
	util "github.com/xxorde/pgglaskugel/util"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

const (
	// Number of bytes to read per iteration
	nBytes = 64
)

var (

	// WaitGroup for workers
	wg sync.WaitGroup

	basebackupCmd = &cobra.Command{
		Use:   "basebackup",
		Short: "Creates a new basebackup from the database",
		Long:  `Creates a new basebackup from the database with the given method.`,
		Run: func(cmd *cobra.Command, args []string) {
			log.Info("Perform basebackup")
			// Get time, name and path for basebackup
			backupTime := startTime.Format(util.BackupTimeFormat)
			backupName := "bb@" + backupTime + ".zst"
			conString := "'" + viper.GetString("connection") + "'"

			// Command to use pg_basebackup
			// Tar format, set backupName as label, make fast checkpoints, return output on standardout
			backupCmd := exec.Command("pg_basebackup", "--pgdata", conString, "--format=tar", "--label", backupName, "--checkpoint", "fast", "--pgdata", "-")
			if viper.GetBool("standalone") {
				// Set command to include WAL files
				backupCmd = exec.Command("pg_basebackup", "--pgdata", conString, "--format=tar", "--label", backupName, "--checkpoint", "fast", "--pgdata", "-", "--xlog-method=fetch")
			}
			log.Debug("backupCmd: ", backupCmd)

			// attach pipe to the command
			backupStdout, err := backupCmd.StdoutPipe()
			if err != nil {
				log.Fatal("Can not attach pipe to backup process, ", err)
			}

			// Watch output on stderror
			backupStderror, err := backupCmd.StderrPipe()
			ec.Check(err)
			go util.WatchOutput(backupStderror, log.Info)

			// This command is used to take the backup and compress it
			compressCmd := exec.Command("zstd")

			// attach pipe to the command
			compressStdout, err := compressCmd.StdoutPipe()
			if err != nil {
				log.Fatal("Can not attach pipe to backup process, ", err)
			}

			// Watch output on stderror
			compressStderror, err := compressCmd.StderrPipe()
			ec.Check(err)
			go util.WatchOutput(compressStderror, log.Info)

			// Pipe the backup in the compression
			compressCmd.Stdin = backupStdout

			// Start the process (in the background)
			if err := backupCmd.Start(); err != nil {
				log.Fatal("pg_basebackup failed on startup, ", err)
			}
			log.Info("Backup was started")

			// Start backup and compression
			if err := compressCmd.Start(); err != nil {
				log.Fatal("zstd failed on startup, ", err)
			}
			log.Info("Compression started")

			// Start worker
			// Add one worker to our waiting group (for waiting later)
			wg.Add(1)
			go handleBackupStream(compressStdout, backupName, &wg)

			// Wait for workers to finish
			//(WAIT FIRST FOR THE WORKER OR WE CAN LOOSE DATA)
			wg.Wait()

			// Wait for backup to finish
			// If there is still data in the output pipe it can be lost!
			err = backupCmd.Wait()
			if err != nil {
				log.Fatal("pg_basebackup failed after startup, ", err)
			}
			log.Debug("backupCmd done")

			// Wait for compression to finish
			// If there is still data in the output pipe it can be lost!
			err = compressCmd.Wait()
			if err != nil {
				log.Fatal("compression failed after startup, ", err)
			}
			log.Debug("compressCmd done")
			printDone()
		},
	}
)

// handleBackupStream takes a stream and persists it with the configured method
func handleBackupStream(input io.ReadCloser, filename string, wg *sync.WaitGroup) {
	// Tell the waiting group this process is done when function ends
	defer wg.Done()

	backupTo := viper.GetString("backup_to")
	switch backupTo {
	case "file":
		writeStreamToFile(input, filename)
	case "s3":
		writeStreamToS3(input, filename)
	default:
		log.Fatal(backupTo, " no valid value for backupTo")
	}
}

// writeStreamToFile handles a stream and writes it to a local file
func writeStreamToFile(input io.ReadCloser, backupName string) {
	backupPath := viper.GetString("archivedir") + "/basebackup/" + backupName

	file, err := os.OpenFile(backupPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0660)
	if err != nil {
		log.Fatal("Can not create output file, ", err)
	}
	defer file.Close()

	log.Debug("Start writing to file")
	written, err := io.Copy(file, input)
	if err != nil {
		log.Fatalf("writeStreamToFile: Error while writing to %s, written %d, error: %v", backupPath, written, err)
	}

	log.Infof("%d bytes were written, waiting for file.Sync()", written)
	file.Sync()
}

// writeStreamToS3 handles a stream and writes it to S3 storage
func writeStreamToS3(input io.ReadCloser, backupName string) {
	bucket := viper.GetString("s3_bucket_backup")
	location := viper.GetString("s3_location")

	// Initialize minio client object.
	minioClient := getS3Connection()

	// Test if bucket is there
	exists, err := minioClient.BucketExists(bucket)
	if err != nil {
		log.Fatal(err)
	}
	if exists {
		log.Debugf("Bucket already exists, we are using it: %s", bucket)
	} else {
		// Try to create bucket
		err = minioClient.MakeBucket(bucket, location)
		if err != nil {
			log.Fatal(err)
		}
		log.Infof("Bucket %s created.", bucket)
	}
	n, err := minioClient.PutObject(bucket, backupName, input, "")
	if err != nil {
		log.Fatal(err)
		return
	}

	log.Infof("Written %d bytes to %s in bucket %s.", n, backupName, bucket)
}

func init() {
	RootCmd.AddCommand(basebackupCmd)
	RootCmd.PersistentFlags().Bool("standalone", false, "Include WAL files in backup so it can be used stand alone")
	// Bind flags to viper
	viper.BindPFlag("standalone", RootCmd.PersistentFlags().Lookup("standalone"))
}
<commit_msg>Add encryption to basebackup<commit_after>// Copyright © 2017 Alexander Sosna <alexander@xxor.de>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cmd

import (
	"io"
	"os"
	"os/exec"
	"sync"

	log "github.com/Sirupsen/logrus"
	ec "github.com/xxorde/pgglaskugel/errorcheck"
	util "github.com/xxorde/pgglaskugel/util"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

const (
	// Number of bytes to read per iteration
	nBytes = 64
)

var (

	// WaitGroup for workers
	wg sync.WaitGroup

	basebackupCmd = &cobra.Command{
		Use:   "basebackup",
		Short: "Creates a new basebackup from the database",
		Long:  `Creates a new basebackup from the database with the given method.`,
		Run: func(cmd *cobra.Command, args []string) {
			log.Info("Perform basebackup")
			// Get time, name and path for basebackup
			backupTime := startTime.Format(util.BackupTimeFormat)
			backupName := "bb@" + backupTime + ".zst"
			conString := viper.GetString("connection")

			log.Debug("conString: ", conString)

			// Command to use pg_basebackup
			// Tar format, set backupName as label, make fast checkpoints, return output on standardout
			backupCmd := exec.Command("pg_basebackup", "--dbname", conString, "--format=tar", "--label", backupName, "--checkpoint", "fast", "--pgdata", "-")
			if viper.GetBool("standalone") {
				// Set command to include WAL files
				backupCmd = exec.Command("pg_basebackup", "--dbname", conString, "--format=tar", "--label", backupName, "--checkpoint", "fast", "--pgdata", "-", "--xlog-method=fetch")
			}
			log.Debug("backupCmd: ", backupCmd)

			// attach pipe to the command
			backupStdout, err := backupCmd.StdoutPipe()
			if err != nil {
				log.Fatal("Can not attach pipe to backup process, ", err)
			}

			// Watch output on stderror
			backupStderror, err := backupCmd.StderrPipe()
			ec.Check(err)
			go util.WatchOutput(backupStderror, log.Info)

			// This command is used to take the backup and compress it
			compressCmd := exec.Command("zstd")

			// attach pipe to the command
			compressStdout, err := compressCmd.StdoutPipe()
			if err != nil {
				log.Fatal("Can not attach pipe to backup process, ", err)
			}

			// Watch output on stderror
			compressStderror, err := compressCmd.StderrPipe()
			ec.Check(err)
			go util.WatchOutput(compressStderror, log.Info)

			// Pipe the backup in the compression
			compressCmd.Stdin = backupStdout

			// Start the process (in the background)
			if err := backupCmd.Start(); err != nil {
				log.Fatal("pg_basebackup failed on startup, ", err)
			}
			log.Info("Backup was started")

			// Start backup and compression
			if err := compressCmd.Start(); err != nil {
				log.Fatal("zstd failed on startup, ", err)
			}
			log.Info("Compression started")

			// Start worker
			// Add one worker to our waiting group (for waiting later)
			wg.Add(1)
			go handleBackupStream(compressStdout, backupName, &wg)

			// Wait for workers to finish
			//(WAIT FIRST FOR THE WORKER OR WE CAN LOOSE DATA)
			wg.Wait()

			// Wait for backup to finish
			// If there is still data in the output pipe it can be lost!
			err = backupCmd.Wait()
			if err != nil {
				log.Fatal("pg_basebackup failed after startup, ", err)
			}
			log.Debug("backupCmd done")

			// Wait for compression to finish
			// If there is still data in the output pipe it can be lost!
			err = compressCmd.Wait()
			if err != nil {
				log.Fatal("compression failed after startup, ", err)
			}
			log.Debug("compressCmd done")
			printDone()
		},
	}
)

// handleBackupStream takes a stream and persists it with the configured method
func handleBackupStream(input io.ReadCloser, filename string, wg *sync.WaitGroup) {
	// Tell the waiting group this process is done when function ends
	defer wg.Done()

	backupTo := viper.GetString("backup_to")
	switch backupTo {
	case "file":
		writeStreamToFile(input, filename)
	case "s3":
		writeStreamToS3(input, filename)
	default:
		log.Fatal(backupTo, " no valid value for backupTo")
	}
}

// writeStreamToFile handles a stream and writes it to a local file
func writeStreamToFile(input io.ReadCloser, backupName string) {
	backupPath := viper.GetString("archivedir") + "/basebackup/" + backupName

	file, err := os.OpenFile(backupPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0660)
	if err != nil {
		log.Fatal("Can not create output file, ", err)
	}
	defer file.Close()

	log.Debug("Start writing to file")
	written, err := io.Copy(file, input)
	if err != nil {
		log.Fatalf("writeStreamToFile: Error while writing to %s, written %d, error: %v", backupPath, written, err)
	}

	log.Infof("%d bytes were written, waiting for file.Sync()", written)
	file.Sync()
}

// writeStreamToS3 handles a stream and writes it to S3 storage
func writeStreamToS3(input io.ReadCloser, backupName string) {
	bucket := viper.GetString("s3_bucket_backup")
	location := viper.GetString("s3_location")
	encrypt := viper.GetBool("encrypt")
	recipient := viper.GetString("recipient")
	contentType := "pgBasebackup"

	var s3Input io.ReadCloser

	// Variables need for encryption
	var gpgCmd *exec.Cmd
	if encrypt {
		log.Debug("Encrypt data, encrypt: ", encrypt)
		// Encrypt the compressed data
		gpgCmd = exec.Command(cmdGpg, "--encrypt", "-o", "-", "--recipient", recipient)
		// Set the encryption output as input for S3
		var err error
		s3Input, err = gpgCmd.StdoutPipe()
		if err != nil {
			log.Fatal("Can not attach pipe to gpg process, ", err)
		}
		// Attach output of WAL to stdin
		gpgCmd.Stdin = input
		// Watch output on stderror
		gpgStderror, err := gpgCmd.StderrPipe()
		ec.Check(err)
		go util.WatchOutput(gpgStderror, log.Warn)

		// Start encryption
		if err := gpgCmd.Start(); err != nil {
			log.Fatal("gpg failed on startup, ", err)
		}
		log.Debug("gpg started")
		contentType = "pgp"
	} else {
		s3Input = input
	}

	// Initialize minio client object.
	minioClient := getS3Connection()

	// Test if bucket is there
	exists, err := minioClient.BucketExists(bucket)
	if err != nil {
		log.Fatal(err)
	}
	if exists {
		log.Debugf("Bucket already exists, we are using it: %s", bucket)
	} else {
		// Try to create bucket
		err = minioClient.MakeBucket(bucket, location)
		if err != nil {
			log.Fatal(err)
		}
		log.Infof("Bucket %s created.", bucket)
	}
	n, err := minioClient.PutObject(bucket, backupName, s3Input, contentType)
	if err != nil {
		log.Fatal(err)
		return
	}
	log.Infof("Written %d bytes to %s in bucket %s.", n, backupName, bucket)

	if encrypt {
		err = gpgCmd.Wait()
		if err != nil {
			log.Fatal("gpg failed after startup, ", err)
		} else {
			log.Debug("Encryption done")
		}
	}
}

func init() {
	RootCmd.AddCommand(basebackupCmd)
	RootCmd.PersistentFlags().Bool("standalone", false, "Include WAL files in backup so it can be used stand alone")
	// Bind flags to viper
	viper.BindPFlag("standalone", RootCmd.PersistentFlags().Lookup("standalone"))
}
<|endoftext|>


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

package main

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/token"
	"go/types"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"unicode"
	"unicode/utf8"

	"golang.org/x/mobile/bind"
	"golang.org/x/mobile/internal/importers"
	"golang.org/x/mobile/internal/importers/java"
	"golang.org/x/mobile/internal/importers/objc"
	"golang.org/x/tools/go/packages"
)

func genPkg(lang string, p *types.Package, astFiles []*ast.File, allPkg []*types.Package, classes []*java.Class, otypes []*objc.Named) {
	fname := defaultFileName(lang, p)
	conf := &bind.GeneratorConfig{
		Fset:   fset,
		Pkg:    p,
		AllPkg: allPkg,
	}
	var pname string
	if p != nil {
		pname = p.Name()
	} else {
		pname = "universe"
	}
	var buf bytes.Buffer
	generator := &bind.Generator{
		Printer: &bind.Printer{Buf: &buf, IndentEach: []byte("\t")},
		Fset:    conf.Fset,
		AllPkg:  conf.AllPkg,
		Pkg:     conf.Pkg,
		Files:   astFiles,
	}
	switch lang {
	case "java":
		g := &bind.JavaGen{
			JavaPkg:   *javaPkg,
			Generator: generator,
		}
		g.Init(classes)

		pkgname := bind.JavaPkgName(*javaPkg, p)
		pkgDir := strings.Replace(pkgname, ".", "/", -1)
		buf.Reset()
		w, closer := writer(filepath.Join("java", pkgDir, fname))
		processErr(g.GenJava())
		io.Copy(w, &buf)
		closer()
		for i, name := range g.ClassNames() {
			buf.Reset()
			w, closer := writer(filepath.Join("java", pkgDir, name+".java"))
			processErr(g.GenClass(i))
			io.Copy(w, &buf)
			closer()
		}
		buf.Reset()
		w, closer = writer(filepath.Join("src", "gobind", pname+"_android.c"))
		processErr(g.GenC())
		io.Copy(w, &buf)
		closer()
		buf.Reset()
		w, closer = writer(filepath.Join("src", "gobind", pname+"_android.h"))
		processErr(g.GenH())
		io.Copy(w, &buf)
		closer()
		// Generate support files along with the universe package
		if p == nil {
			dir, err := packageDir("golang.org/x/mobile/bind")
			if err != nil {
				errorf(`"golang.org/x/mobile/bind" is not found; run go get golang.org/x/mobile/bind: %v`, err)
				return
			}
			repo := filepath.Clean(filepath.Join(dir, "..")) // golang.org/x/mobile directory.
			for _, javaFile := range []string{"Seq.java"} {
				src := filepath.Join(repo, "bind/java/"+javaFile)
				in, err := os.Open(src)
				if err != nil {
					errorf("failed to open Java support file: %v", err)
				}
				defer in.Close()
				w, closer := writer(filepath.Join("java", "go", javaFile))
				defer closer()
				if _, err := io.Copy(w, in); err != nil {
					errorf("failed to copy Java support file: %v", err)
					return
				}
			}
			// Copy support files
			if err != nil {
				errorf("unable to import bind/java: %v", err)
				return
			}
			javaDir, err := packageDir("golang.org/x/mobile/bind/java")
			if err != nil {
				errorf("unable to import bind/java: %v", err)
				return
			}
			copyFile(filepath.Join("src", "gobind", "seq_android.c"), filepath.Join(javaDir, "seq_android.c.support"))
			copyFile(filepath.Join("src", "gobind", "seq_android.go"), filepath.Join(javaDir, "seq_android.go.support"))
			copyFile(filepath.Join("src", "gobind", "seq_android.h"), filepath.Join(javaDir, "seq_android.h"))
		}
	case "go":
		w, closer := writer(filepath.Join("src", "gobind", fname))
		conf.Writer = w
		processErr(bind.GenGo(conf))
		closer()
		w, closer = writer(filepath.Join("src", "gobind", pname+".h"))
		genPkgH(w, pname)
		io.Copy(w, &buf)
		closer()
		w, closer = writer(filepath.Join("src", "gobind", "seq.h"))
		genPkgH(w, "seq")
		io.Copy(w, &buf)
		closer()
		dir, err := packageDir("golang.org/x/mobile/bind")
		if err != nil {
			errorf("unable to import bind: %v", err)
			return
		}
		copyFile(filepath.Join("src", "gobind", "seq.go"), filepath.Join(dir, "seq.go.support"))
	case "objc":
		g := &bind.ObjcGen{
			Generator: generator,
			Prefix:    *prefix,
		}
		g.Init(otypes)
		w, closer := writer(filepath.Join("src", "gobind", pname+"_darwin.h"))
		processErr(g.GenGoH())
		io.Copy(w, &buf)
		closer()
		hname := strings.Title(fname[:len(fname)-2]) + ".objc.h"
		w, closer = writer(filepath.Join("src", "gobind", hname))
		processErr(g.GenH())
		io.Copy(w, &buf)
		closer()
		mname := strings.Title(fname[:len(fname)-2]) + "_darwin.m"
		w, closer = writer(filepath.Join("src", "gobind", mname))
		conf.Writer = w
		processErr(g.GenM())
		io.Copy(w, &buf)
		closer()
		if p == nil {
			// Copy support files
			dir, err := packageDir("golang.org/x/mobile/bind/objc")
			if err != nil {
				errorf("unable to import bind/objc: %v", err)
				return
			}
			copyFile(filepath.Join("src", "gobind", "seq_darwin.m"), filepath.Join(dir, "seq_darwin.m.support"))
			copyFile(filepath.Join("src", "gobind", "seq_darwin.go"), filepath.Join(dir, "seq_darwin.go.support"))
			copyFile(filepath.Join("src", "gobind", "ref.h"), filepath.Join(dir, "ref.h"))
			copyFile(filepath.Join("src", "gobind", "seq_darwin.h"), filepath.Join(dir, "seq_darwin.h"))
		}
	default:
		errorf("unknown target language: %q", lang)
	}
}

func genPkgH(w io.Writer, pname string) {
	fmt.Fprintf(w, `// Code generated by gobind. DO NOT EDIT.

#ifdef __GOBIND_ANDROID__
#include "%[1]s_android.h"
#endif
#ifdef __GOBIND_DARWIN__
#include "%[1]s_darwin.h"
#endif`, pname)
}

func genObjcPackages(dir string, types []*objc.Named, embedders []importers.Struct) error {
	var buf bytes.Buffer
	cg := &bind.ObjcWrapper{
		Printer: &bind.Printer{
			IndentEach: []byte("\t"),
			Buf:        &buf,
		},
	}
	var genNames []string
	for _, emb := range embedders {
		genNames = append(genNames, emb.Name)
	}
	cg.Init(types, genNames)
	for i, opkg := range cg.Packages() {
		pkgDir := filepath.Join(dir, "src", "ObjC", opkg)
		if err := os.MkdirAll(pkgDir, 0700); err != nil {
			return err
		}
		pkgFile := filepath.Join(pkgDir, "package.go")
		buf.Reset()
		cg.GenPackage(i)
		if err := ioutil.WriteFile(pkgFile, buf.Bytes(), 0600); err != nil {
			return err
		}
	}
	buf.Reset()
	cg.GenInterfaces()
	objcBase := filepath.Join(dir, "src", "ObjC")
	if err := os.MkdirAll(objcBase, 0700); err != nil {
		return err
	}
	if err := ioutil.WriteFile(filepath.Join(objcBase, "interfaces.go"), buf.Bytes(), 0600); err != nil {
		return err
	}
	goBase := filepath.Join(dir, "src", "gobind")
	if err := os.MkdirAll(goBase, 0700); err != nil {
		return err
	}
	buf.Reset()
	cg.GenGo()
	if err := ioutil.WriteFile(filepath.Join(goBase, "interfaces_darwin.go"), buf.Bytes(), 0600); err != nil {
		return err
	}
	buf.Reset()
	cg.GenH()
	if err := ioutil.WriteFile(filepath.Join(goBase, "interfaces.h"), buf.Bytes(), 0600); err != nil {
		return err
	}
	buf.Reset()
	cg.GenM()
	if err := ioutil.WriteFile(filepath.Join(goBase, "interfaces_darwin.m"), buf.Bytes(), 0600); err != nil {
		return err
	}
	return nil
}

func genJavaPackages(dir string, classes []*java.Class, embedders []importers.Struct) error {
	var buf bytes.Buffer
	cg := &bind.ClassGen{
		JavaPkg: *javaPkg,
		Printer: &bind.Printer{
			IndentEach: []byte("\t"),
			Buf:        &buf,
		},
	}
	cg.Init(classes, embedders)
	for i, jpkg := range cg.Packages() {
		pkgDir := filepath.Join(dir, "src", "Java", jpkg)
		if err := os.MkdirAll(pkgDir, 0700); err != nil {
			return err
		}
		pkgFile := filepath.Join(pkgDir, "package.go")
		buf.Reset()
		cg.GenPackage(i)
		if err := ioutil.WriteFile(pkgFile, buf.Bytes(), 0600); err != nil {
			return err
		}
	}
	buf.Reset()
	cg.GenInterfaces()
	javaBase := filepath.Join(dir, "src", "Java")
	if err := os.MkdirAll(javaBase, 0700); err != nil {
		return err
	}
	if err := ioutil.WriteFile(filepath.Join(javaBase, "interfaces.go"), buf.Bytes(), 0600); err != nil {
		return err
	}
	goBase := filepath.Join(dir, "src", "gobind")
	if err := os.MkdirAll(goBase, 0700); err != nil {
		return err
	}
	buf.Reset()
	cg.GenGo()
	if err := ioutil.WriteFile(filepath.Join(goBase, "classes_android.go"), buf.Bytes(), 0600); err != nil {
		return err
	}
	buf.Reset()
	cg.GenH()
	if err := ioutil.WriteFile(filepath.Join(goBase, "classes.h"), buf.Bytes(), 0600); err != nil {
		return err
	}
	buf.Reset()
	cg.GenC()
	if err := ioutil.WriteFile(filepath.Join(goBase, "classes_android.c"), buf.Bytes(), 0600); err != nil {
		return err
	}
	return nil
}

func processErr(err error) {
	if err != nil {
		if list, _ := err.(bind.ErrorList); len(list) > 0 {
			for _, err := range list {
				errorf("%v", err)
			}
		} else {
			errorf("%v", err)
		}
	}
}

var fset = token.NewFileSet()

func writer(fname string) (w io.Writer, closer func()) {
	if *outdir == "" {
		return os.Stdout, func() { return }
	}

	name := filepath.Join(*outdir, fname)
	dir := filepath.Dir(name)
	if err := os.MkdirAll(dir, 0755); err != nil {
		errorf("invalid output dir: %v", err)
		os.Exit(exitStatus)
	}

	f, err := os.Create(name)
	if err != nil {
		errorf("invalid output dir: %v", err)
		os.Exit(exitStatus)
	}
	closer = func() {
		if err := f.Close(); err != nil {
			errorf("error in closing output file: %v", err)
		}
	}
	return f, closer
}

func copyFile(dst, src string) {
	w, closer := writer(dst)
	f, err := os.Open(src)
	if err != nil {
		errorf("unable to open file: %v", err)
		closer()
		os.Exit(exitStatus)
	}
	if _, err := io.Copy(w, f); err != nil {
		errorf("unable to copy file: %v", err)
		f.Close()
		closer()
		os.Exit(exitStatus)
	}
	f.Close()
	closer()
}

func defaultFileName(lang string, pkg *types.Package) string {
	switch lang {
	case "java":
		if pkg == nil {
			return "Universe.java"
		}
		firstRune, size := utf8.DecodeRuneInString(pkg.Name())
		className := string(unicode.ToUpper(firstRune)) + pkg.Name()[size:]
		return className + ".java"
	case "go":
		if pkg == nil {
			return "go_main.go"
		}
		return "go_" + pkg.Name() + "main.go"
	case "objc":
		if pkg == nil {
			return "Universe.m"
		}
		firstRune, size := utf8.DecodeRuneInString(pkg.Name())
		className := string(unicode.ToUpper(firstRune)) + pkg.Name()[size:]
		return *prefix + className + ".m"
	}
	errorf("unknown target language: %q", lang)
	os.Exit(exitStatus)
	return ""
}

func packageDir(path string) (string, error) {
	pkgs, err := packages.Load(nil, path)
	if err != nil {
		return "", err
	}
	pkg := pkgs[0]
	if len(pkg.Errors) > 0 {
		return "", fmt.Errorf("%v", pkg.Errors)
	}
	return filepath.Dir(pkg.GoFiles[0]), nil
}
<commit_msg>cmd/gobind: do not compile package just to find package dir<commit_after>// Copyright 2014 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/token"
	"go/types"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"unicode"
	"unicode/utf8"

	"golang.org/x/mobile/bind"
	"golang.org/x/mobile/internal/importers"
	"golang.org/x/mobile/internal/importers/java"
	"golang.org/x/mobile/internal/importers/objc"
	"golang.org/x/tools/go/packages"
)

func genPkg(lang string, p *types.Package, astFiles []*ast.File, allPkg []*types.Package, classes []*java.Class, otypes []*objc.Named) {
	fname := defaultFileName(lang, p)
	conf := &bind.GeneratorConfig{
		Fset:   fset,
		Pkg:    p,
		AllPkg: allPkg,
	}
	var pname string
	if p != nil {
		pname = p.Name()
	} else {
		pname = "universe"
	}
	var buf bytes.Buffer
	generator := &bind.Generator{
		Printer: &bind.Printer{Buf: &buf, IndentEach: []byte("\t")},
		Fset:    conf.Fset,
		AllPkg:  conf.AllPkg,
		Pkg:     conf.Pkg,
		Files:   astFiles,
	}
	switch lang {
	case "java":
		g := &bind.JavaGen{
			JavaPkg:   *javaPkg,
			Generator: generator,
		}
		g.Init(classes)

		pkgname := bind.JavaPkgName(*javaPkg, p)
		pkgDir := strings.Replace(pkgname, ".", "/", -1)
		buf.Reset()
		w, closer := writer(filepath.Join("java", pkgDir, fname))
		processErr(g.GenJava())
		io.Copy(w, &buf)
		closer()
		for i, name := range g.ClassNames() {
			buf.Reset()
			w, closer := writer(filepath.Join("java", pkgDir, name+".java"))
			processErr(g.GenClass(i))
			io.Copy(w, &buf)
			closer()
		}
		buf.Reset()
		w, closer = writer(filepath.Join("src", "gobind", pname+"_android.c"))
		processErr(g.GenC())
		io.Copy(w, &buf)
		closer()
		buf.Reset()
		w, closer = writer(filepath.Join("src", "gobind", pname+"_android.h"))
		processErr(g.GenH())
		io.Copy(w, &buf)
		closer()
		// Generate support files along with the universe package
		if p == nil {
			dir, err := packageDir("golang.org/x/mobile/bind")
			if err != nil {
				errorf(`"golang.org/x/mobile/bind" is not found; run go get golang.org/x/mobile/bind: %v`, err)
				return
			}
			repo := filepath.Clean(filepath.Join(dir, "..")) // golang.org/x/mobile directory.
			for _, javaFile := range []string{"Seq.java"} {
				src := filepath.Join(repo, "bind/java/"+javaFile)
				in, err := os.Open(src)
				if err != nil {
					errorf("failed to open Java support file: %v", err)
				}
				defer in.Close()
				w, closer := writer(filepath.Join("java", "go", javaFile))
				defer closer()
				if _, err := io.Copy(w, in); err != nil {
					errorf("failed to copy Java support file: %v", err)
					return
				}
			}
			// Copy support files
			if err != nil {
				errorf("unable to import bind/java: %v", err)
				return
			}
			javaDir, err := packageDir("golang.org/x/mobile/bind/java")
			if err != nil {
				errorf("unable to import bind/java: %v", err)
				return
			}
			copyFile(filepath.Join("src", "gobind", "seq_android.c"), filepath.Join(javaDir, "seq_android.c.support"))
			copyFile(filepath.Join("src", "gobind", "seq_android.go"), filepath.Join(javaDir, "seq_android.go.support"))
			copyFile(filepath.Join("src", "gobind", "seq_android.h"), filepath.Join(javaDir, "seq_android.h"))
		}
	case "go":
		w, closer := writer(filepath.Join("src", "gobind", fname))
		conf.Writer = w
		processErr(bind.GenGo(conf))
		closer()
		w, closer = writer(filepath.Join("src", "gobind", pname+".h"))
		genPkgH(w, pname)
		io.Copy(w, &buf)
		closer()
		w, closer = writer(filepath.Join("src", "gobind", "seq.h"))
		genPkgH(w, "seq")
		io.Copy(w, &buf)
		closer()
		dir, err := packageDir("golang.org/x/mobile/bind")
		if err != nil {
			errorf("unable to import bind: %v", err)
			return
		}
		copyFile(filepath.Join("src", "gobind", "seq.go"), filepath.Join(dir, "seq.go.support"))
	case "objc":
		g := &bind.ObjcGen{
			Generator: generator,
			Prefix:    *prefix,
		}
		g.Init(otypes)
		w, closer := writer(filepath.Join("src", "gobind", pname+"_darwin.h"))
		processErr(g.GenGoH())
		io.Copy(w, &buf)
		closer()
		hname := strings.Title(fname[:len(fname)-2]) + ".objc.h"
		w, closer = writer(filepath.Join("src", "gobind", hname))
		processErr(g.GenH())
		io.Copy(w, &buf)
		closer()
		mname := strings.Title(fname[:len(fname)-2]) + "_darwin.m"
		w, closer = writer(filepath.Join("src", "gobind", mname))
		conf.Writer = w
		processErr(g.GenM())
		io.Copy(w, &buf)
		closer()
		if p == nil {
			// Copy support files
			dir, err := packageDir("golang.org/x/mobile/bind/objc")
			if err != nil {
				errorf("unable to import bind/objc: %v", err)
				return
			}
			copyFile(filepath.Join("src", "gobind", "seq_darwin.m"), filepath.Join(dir, "seq_darwin.m.support"))
			copyFile(filepath.Join("src", "gobind", "seq_darwin.go"), filepath.Join(dir, "seq_darwin.go.support"))
			copyFile(filepath.Join("src", "gobind", "ref.h"), filepath.Join(dir, "ref.h"))
			copyFile(filepath.Join("src", "gobind", "seq_darwin.h"), filepath.Join(dir, "seq_darwin.h"))
		}
	default:
		errorf("unknown target language: %q", lang)
	}
}

func genPkgH(w io.Writer, pname string) {
	fmt.Fprintf(w, `// Code generated by gobind. DO NOT EDIT.

#ifdef __GOBIND_ANDROID__
#include "%[1]s_android.h"
#endif
#ifdef __GOBIND_DARWIN__
#include "%[1]s_darwin.h"
#endif`, pname)
}

func genObjcPackages(dir string, types []*objc.Named, embedders []importers.Struct) error {
	var buf bytes.Buffer
	cg := &bind.ObjcWrapper{
		Printer: &bind.Printer{
			IndentEach: []byte("\t"),
			Buf:        &buf,
		},
	}
	var genNames []string
	for _, emb := range embedders {
		genNames = append(genNames, emb.Name)
	}
	cg.Init(types, genNames)
	for i, opkg := range cg.Packages() {
		pkgDir := filepath.Join(dir, "src", "ObjC", opkg)
		if err := os.MkdirAll(pkgDir, 0700); err != nil {
			return err
		}
		pkgFile := filepath.Join(pkgDir, "package.go")
		buf.Reset()
		cg.GenPackage(i)
		if err := ioutil.WriteFile(pkgFile, buf.Bytes(), 0600); err != nil {
			return err
		}
	}
	buf.Reset()
	cg.GenInterfaces()
	objcBase := filepath.Join(dir, "src", "ObjC")
	if err := os.MkdirAll(objcBase, 0700); err != nil {
		return err
	}
	if err := ioutil.WriteFile(filepath.Join(objcBase, "interfaces.go"), buf.Bytes(), 0600); err != nil {
		return err
	}
	goBase := filepath.Join(dir, "src", "gobind")
	if err := os.MkdirAll(goBase, 0700); err != nil {
		return err
	}
	buf.Reset()
	cg.GenGo()
	if err := ioutil.WriteFile(filepath.Join(goBase, "interfaces_darwin.go"), buf.Bytes(), 0600); err != nil {
		return err
	}
	buf.Reset()
	cg.GenH()
	if err := ioutil.WriteFile(filepath.Join(goBase, "interfaces.h"), buf.Bytes(), 0600); err != nil {
		return err
	}
	buf.Reset()
	cg.GenM()
	if err := ioutil.WriteFile(filepath.Join(goBase, "interfaces_darwin.m"), buf.Bytes(), 0600); err != nil {
		return err
	}
	return nil
}

func genJavaPackages(dir string, classes []*java.Class, embedders []importers.Struct) error {
	var buf bytes.Buffer
	cg := &bind.ClassGen{
		JavaPkg: *javaPkg,
		Printer: &bind.Printer{
			IndentEach: []byte("\t"),
			Buf:        &buf,
		},
	}
	cg.Init(classes, embedders)
	for i, jpkg := range cg.Packages() {
		pkgDir := filepath.Join(dir, "src", "Java", jpkg)
		if err := os.MkdirAll(pkgDir, 0700); err != nil {
			return err
		}
		pkgFile := filepath.Join(pkgDir, "package.go")
		buf.Reset()
		cg.GenPackage(i)
		if err := ioutil.WriteFile(pkgFile, buf.Bytes(), 0600); err != nil {
			return err
		}
	}
	buf.Reset()
	cg.GenInterfaces()
	javaBase := filepath.Join(dir, "src", "Java")
	if err := os.MkdirAll(javaBase, 0700); err != nil {
		return err
	}
	if err := ioutil.WriteFile(filepath.Join(javaBase, "interfaces.go"), buf.Bytes(), 0600); err != nil {
		return err
	}
	goBase := filepath.Join(dir, "src", "gobind")
	if err := os.MkdirAll(goBase, 0700); err != nil {
		return err
	}
	buf.Reset()
	cg.GenGo()
	if err := ioutil.WriteFile(filepath.Join(goBase, "classes_android.go"), buf.Bytes(), 0600); err != nil {
		return err
	}
	buf.Reset()
	cg.GenH()
	if err := ioutil.WriteFile(filepath.Join(goBase, "classes.h"), buf.Bytes(), 0600); err != nil {
		return err
	}
	buf.Reset()
	cg.GenC()
	if err := ioutil.WriteFile(filepath.Join(goBase, "classes_android.c"), buf.Bytes(), 0600); err != nil {
		return err
	}
	return nil
}

func processErr(err error) {
	if err != nil {
		if list, _ := err.(bind.ErrorList); len(list) > 0 {
			for _, err := range list {
				errorf("%v", err)
			}
		} else {
			errorf("%v", err)
		}
	}
}

var fset = token.NewFileSet()

func writer(fname string) (w io.Writer, closer func()) {
	if *outdir == "" {
		return os.Stdout, func() { return }
	}

	name := filepath.Join(*outdir, fname)
	dir := filepath.Dir(name)
	if err := os.MkdirAll(dir, 0755); err != nil {
		errorf("invalid output dir: %v", err)
		os.Exit(exitStatus)
	}

	f, err := os.Create(name)
	if err != nil {
		errorf("invalid output dir: %v", err)
		os.Exit(exitStatus)
	}
	closer = func() {
		if err := f.Close(); err != nil {
			errorf("error in closing output file: %v", err)
		}
	}
	return f, closer
}

func copyFile(dst, src string) {
	w, closer := writer(dst)
	f, err := os.Open(src)
	if err != nil {
		errorf("unable to open file: %v", err)
		closer()
		os.Exit(exitStatus)
	}
	if _, err := io.Copy(w, f); err != nil {
		errorf("unable to copy file: %v", err)
		f.Close()
		closer()
		os.Exit(exitStatus)
	}
	f.Close()
	closer()
}

func defaultFileName(lang string, pkg *types.Package) string {
	switch lang {
	case "java":
		if pkg == nil {
			return "Universe.java"
		}
		firstRune, size := utf8.DecodeRuneInString(pkg.Name())
		className := string(unicode.ToUpper(firstRune)) + pkg.Name()[size:]
		return className + ".java"
	case "go":
		if pkg == nil {
			return "go_main.go"
		}
		return "go_" + pkg.Name() + "main.go"
	case "objc":
		if pkg == nil {
			return "Universe.m"
		}
		firstRune, size := utf8.DecodeRuneInString(pkg.Name())
		className := string(unicode.ToUpper(firstRune)) + pkg.Name()[size:]
		return *prefix + className + ".m"
	}
	errorf("unknown target language: %q", lang)
	os.Exit(exitStatus)
	return ""
}

func packageDir(path string) (string, error) {
	mode := packages.NeedFiles
	pkgs, err := packages.Load(&packages.Config{Mode: mode}, path)
	if err != nil {
		return "", err
	}
	if len(pkgs) == 0 || len(pkgs[0].GoFiles) == 0 {
		return "", fmt.Errorf("no Go package in %v", path)
	}
	pkg := pkgs[0]
	if len(pkg.Errors) > 0 {
		return "", fmt.Errorf("%v", pkg.Errors)
	}
	return filepath.Dir(pkg.GoFiles[0]), nil
}
<|endoftext|>