<commit_before>// Tor pluggable transports library.
//
// Sample client usage:
//
// PtClientSetup([]string{"foo"})
// ln, err := startSocksListener()
// if err != nil {
// 	panic(err.Error())
// }
// PtCmethod("foo", "socks4", ln.Addr())
// PtCmethodsDone()
//
// Sample server usage:
//
// info := PtServerSetup([]string{"foo", "bar"})
// for _, bindAddr := range info.BindAddrs {
// 	ln, err := startListener(bindAddr.Addr)
// 	if err != nil {
// 		panic(err.Error())
// 	}
// 	PtSmethod(bindAddr.MethodName, ln.Addr())
// }
// PtSmethodsDone()

package main

import (
	"bytes"
	"fmt"
	"net"
	"os"
	"strings"
)

// Abort with an ENV-ERROR if the environment variable isn't set.
func getenvRequired(key string) string {
	value := os.Getenv(key)
	if value == "" {
		PtEnvError(fmt.Sprintf("no %s environment variable", key))
	}
	return value
}

// Escape a string so it contains no byte values over 127 and doesn't contain
// any of the characters '\x00', '\n', or '\\'.
func escape(s string) string {
	var buf bytes.Buffer
	for _, b := range []byte(s) {
		if b == '\n' {
			buf.WriteString("\\n")
		} else if b == '\\' {
			buf.WriteString("\\\\")
		} else if 0 < b && b < 128 {
			buf.WriteByte(b)
		} else {
			fmt.Fprintf(&buf, "\\x%02x", b)
		}
	}
	return buf.String()
}

// Print a pluggable transports protocol line to stdout. The line consists of an
// unescaped keyword, followed by any number of escaped strings.
func PtLine(keyword string, v ...string) {
	var buf bytes.Buffer
	buf.WriteString(keyword)
	for _, x := range v {
		buf.WriteString(" " + escape(x))
	}
	fmt.Println(buf.String())
}

// All of the Pt*Error functions call os.Exit(1).

// Emit an ENV-ERROR with explanation text.
func PtEnvError(msg string) {
	PtLine("ENV-ERROR", msg)
	os.Exit(1)
}

// Emit a VERSION-ERROR with explanation text.
func PtVersionError(msg string) {
	PtLine("VERSION-ERROR", msg)
	os.Exit(1)
}

// Emit a CMETHOD-ERROR with explanation text.
func PtCmethodError(methodName, msg string) {
	PtLine("CMETHOD-ERROR", methodName, msg)
	os.Exit(1)
}

// Emit an SMETHOD-ERROR with explanation text.
func PtSmethodError(methodName, msg string) {
	PtLine("SMETHOD-ERROR", methodName, msg)
	os.Exit(1)
}

// Emit a CMETHOD line. socks must be "socks4" or "socks5". Call this once for
// each listening client SOCKS port.
func PtCmethod(name string, socks string, addr net.Addr) {
	PtLine("CMETHOD", name, socks, addr.String())
}

// Emit a CMETHODS DONE line. Call this after opening all client listeners.
func PtCmethodsDone() {
	PtLine("CMETHODS", "DONE")
}

// Emit an SMETHOD line. Call this once for each listening server port.
func PtSmethod(name string, addr net.Addr) {
	PtLine("SMETHOD", name, addr.String())
}

// Emit an SMETHODS DONE line. Call this after opening all server listeners.
func PtSmethodsDone() {
	PtLine("SMETHODS", "DONE")
}

// Get a pluggable transports version offered by Tor and understood by us, if
// any. The only version we understand is "1". This function reads the
// environment variable TOR_PT_MANAGED_TRANSPORT_VER.
func getManagedTransportVer() string {
	const transportVersion = "1"
	for _, offered := range strings.Split(getenvRequired("TOR_PT_MANAGED_TRANSPORT_VER"), ",") {
		if offered == transportVersion {
			return offered
		}
	}
	return ""
}

// Get the intersection of the method names offered by Tor and those in
// methodNames. This function reads the environment variable
// TOR_PT_CLIENT_TRANSPORTS.
func getClientTransports(methodNames []string) []string {
	clientTransports := getenvRequired("TOR_PT_CLIENT_TRANSPORTS")
	if clientTransports == "*" {
		return methodNames
	}
	result := make([]string, 0)
	for _, requested := range strings.Split(clientTransports, ",") {
		for _, methodName := range methodNames {
			if requested == methodName {
				result = append(result, methodName)
				break
			}
		}
	}
	return result
}

// This structure is returned by PtClientSetup. It consists of a list of method
// names.
type PtClientInfo struct {
	MethodNames []string
}

// Check the client pluggable transports environments, emitting an error message
// and exiting the program if any error is encountered. Returns a subset of
// methodNames requested by Tor.
func PtClientSetup(methodNames []string) PtClientInfo {
	var info PtClientInfo

	ver := getManagedTransportVer()
	if ver == "" {
		PtVersionError("no-version")
	} else {
		PtLine("VERSION", ver)
	}

	info.MethodNames = getClientTransports(methodNames)
	if len(info.MethodNames) == 0 {
		PtCmethodsDone()
		os.Exit(1)
	}

	return info
}

// A combination of a method name and an address, as extracted from
// TOR_PT_SERVER_BINDADDR.
type PtBindAddr struct {
	MethodName string
	Addr       *net.TCPAddr
}

// Resolve an address string into a net.TCPAddr.
func resolveBindAddr(bindAddr string) (*net.TCPAddr, error) {
	addr, err := net.ResolveTCPAddr("tcp", bindAddr)
	if err == nil {
		return addr, nil
	}
	// Before the fixing of bug #7011, tor doesn't put brackets around IPv6
	// addresses. Split after the last colon, assuming it is a port
	// separator, and try adding the brackets.
	parts := strings.Split(bindAddr, ":")
	if len(parts) <= 2 {
		return nil, err
	}
	bindAddr = "[" + strings.Join(parts[:len(parts)-1], ":") + "]:" + parts[len(parts)-1]
	return net.ResolveTCPAddr("tcp", bindAddr)
}

// Return a new slice, the members of which are those members of addrs having a
// MethodName in methodsNames.
func filterBindAddrs(addrs []PtBindAddr, methodNames []string) []PtBindAddr {
	var result []PtBindAddr

	for _, ba := range addrs {
		for _, methodName := range methodNames {
			if ba.MethodName == methodName {
				result = append(result, ba)
				break
			}
		}
	}

	return result
}

// Return a map from method names to bind addresses. The map is the contents of
// TOR_PT_SERVER_BINDADDR, with keys filtered by TOR_PT_SERVER_TRANSPORTS, and
// further filtered by the methods in methodNames.
func getServerBindAddrs(methodNames []string) []PtBindAddr {
	var result []PtBindAddr

	// Get the list of all requested bindaddrs.
	var serverBindAddr = getenvRequired("TOR_PT_SERVER_BINDADDR")
	for _, spec := range strings.Split(serverBindAddr, ",") {
		var bindAddr PtBindAddr

		parts := strings.SplitN(spec, "-", 2)
		if len(parts) != 2 {
			PtEnvError(fmt.Sprintf("TOR_PT_SERVER_BINDADDR: %q: doesn't contain \"-\"", spec))
		}
		bindAddr.MethodName = parts[0]
		addr, err := resolveBindAddr(parts[1])
		if err != nil {
			PtEnvError(fmt.Sprintf("TOR_PT_SERVER_BINDADDR: %q: %s", spec, err.Error()))
		}
		bindAddr.Addr = addr
		result = append(result, bindAddr)
	}

	// Filter by TOR_PT_SERVER_TRANSPORTS.
	serverTransports := getenvRequired("TOR_PT_SERVER_TRANSPORTS")
	if serverTransports != "*" {
		result = filterBindAddrs(result, strings.Split(serverTransports, ","))
	}

	// Finally filter by what we understand.
	result = filterBindAddrs(result, methodNames)

	return result
}

// This structure is returned by PtServerSetup. It consists of a list of
// PtBindAddrs, along with a single address for the ORPort.
type PtServerInfo struct {
	BindAddrs []PtBindAddr
	OrAddr    *net.TCPAddr
}

// Check the server pluggable transports environments, emitting an error message
// and exiting the program if any error is encountered. Resolves the various
// requested bind addresses and the server ORPort. Returns a PtServerInfo
// struct.
func PtServerSetup(methodNames []string) PtServerInfo {
	var info PtServerInfo
	var err error

	ver := getManagedTransportVer()
	if ver == "" {
		PtVersionError("no-version")
	} else {
		PtLine("VERSION", ver)
	}

	var orPort = getenvRequired("TOR_PT_ORPORT")
	info.OrAddr, err = net.ResolveTCPAddr("tcp", orPort)
	if err != nil {
		PtEnvError(fmt.Sprintf("cannot resolve TOR_PT_ORPORT %q: %s", orPort, err.Error()))
	}

	info.BindAddrs = getServerBindAddrs(methodNames)
	if len(info.BindAddrs) == 0 {
		PtSmethodsDone()
		os.Exit(1)
	}

	return info
}
<commit_msg>getenv function.<commit_after>// Tor pluggable transports library.
//
// Sample client usage:
//
// PtClientSetup([]string{"foo"})
// ln, err := startSocksListener()
// if err != nil {
// 	panic(err.Error())
// }
// PtCmethod("foo", "socks4", ln.Addr())
// PtCmethodsDone()
//
// Sample server usage:
//
// info := PtServerSetup([]string{"foo", "bar"})
// for _, bindAddr := range info.BindAddrs {
// 	ln, err := startListener(bindAddr.Addr)
// 	if err != nil {
// 		panic(err.Error())
// 	}
// 	PtSmethod(bindAddr.MethodName, ln.Addr())
// }
// PtSmethodsDone()

package main

import (
	"bytes"
	"fmt"
	"net"
	"os"
	"strings"
)

func getenv(key string) string {
	return os.Getenv(key)
}

// Abort with an ENV-ERROR if the environment variable isn't set.
func getenvRequired(key string) string {
	value := os.Getenv(key)
	if value == "" {
		PtEnvError(fmt.Sprintf("no %s environment variable", key))
	}
	return value
}

// Escape a string so it contains no byte values over 127 and doesn't contain
// any of the characters '\x00', '\n', or '\\'.
func escape(s string) string {
	var buf bytes.Buffer
	for _, b := range []byte(s) {
		if b == '\n' {
			buf.WriteString("\\n")
		} else if b == '\\' {
			buf.WriteString("\\\\")
		} else if 0 < b && b < 128 {
			buf.WriteByte(b)
		} else {
			fmt.Fprintf(&buf, "\\x%02x", b)
		}
	}
	return buf.String()
}

// Print a pluggable transports protocol line to stdout. The line consists of an
// unescaped keyword, followed by any number of escaped strings.
func PtLine(keyword string, v ...string) {
	var buf bytes.Buffer
	buf.WriteString(keyword)
	for _, x := range v {
		buf.WriteString(" " + escape(x))
	}
	fmt.Println(buf.String())
}

// All of the Pt*Error functions call os.Exit(1).

// Emit an ENV-ERROR with explanation text.
func PtEnvError(msg string) {
	PtLine("ENV-ERROR", msg)
	os.Exit(1)
}

// Emit a VERSION-ERROR with explanation text.
func PtVersionError(msg string) {
	PtLine("VERSION-ERROR", msg)
	os.Exit(1)
}

// Emit a CMETHOD-ERROR with explanation text.
func PtCmethodError(methodName, msg string) {
	PtLine("CMETHOD-ERROR", methodName, msg)
	os.Exit(1)
}

// Emit an SMETHOD-ERROR with explanation text.
func PtSmethodError(methodName, msg string) {
	PtLine("SMETHOD-ERROR", methodName, msg)
	os.Exit(1)
}

// Emit a CMETHOD line. socks must be "socks4" or "socks5". Call this once for
// each listening client SOCKS port.
func PtCmethod(name string, socks string, addr net.Addr) {
	PtLine("CMETHOD", name, socks, addr.String())
}

// Emit a CMETHODS DONE line. Call this after opening all client listeners.
func PtCmethodsDone() {
	PtLine("CMETHODS", "DONE")
}

// Emit an SMETHOD line. Call this once for each listening server port.
func PtSmethod(name string, addr net.Addr) {
	PtLine("SMETHOD", name, addr.String())
}

// Emit an SMETHODS DONE line. Call this after opening all server listeners.
func PtSmethodsDone() {
	PtLine("SMETHODS", "DONE")
}

// Get a pluggable transports version offered by Tor and understood by us, if
// any. The only version we understand is "1". This function reads the
// environment variable TOR_PT_MANAGED_TRANSPORT_VER.
func getManagedTransportVer() string {
	const transportVersion = "1"
	for _, offered := range strings.Split(getenvRequired("TOR_PT_MANAGED_TRANSPORT_VER"), ",") {
		if offered == transportVersion {
			return offered
		}
	}
	return ""
}

// Get the intersection of the method names offered by Tor and those in
// methodNames. This function reads the environment variable
// TOR_PT_CLIENT_TRANSPORTS.
func getClientTransports(methodNames []string) []string {
	clientTransports := getenvRequired("TOR_PT_CLIENT_TRANSPORTS")
	if clientTransports == "*" {
		return methodNames
	}
	result := make([]string, 0)
	for _, requested := range strings.Split(clientTransports, ",") {
		for _, methodName := range methodNames {
			if requested == methodName {
				result = append(result, methodName)
				break
			}
		}
	}
	return result
}

// This structure is returned by PtClientSetup. It consists of a list of method
// names.
type PtClientInfo struct {
	MethodNames []string
}

// Check the client pluggable transports environments, emitting an error message
// and exiting the program if any error is encountered. Returns a subset of
// methodNames requested by Tor.
func PtClientSetup(methodNames []string) PtClientInfo {
	var info PtClientInfo

	ver := getManagedTransportVer()
	if ver == "" {
		PtVersionError("no-version")
	} else {
		PtLine("VERSION", ver)
	}

	info.MethodNames = getClientTransports(methodNames)
	if len(info.MethodNames) == 0 {
		PtCmethodsDone()
		os.Exit(1)
	}

	return info
}

// A combination of a method name and an address, as extracted from
// TOR_PT_SERVER_BINDADDR.
type PtBindAddr struct {
	MethodName string
	Addr       *net.TCPAddr
}

// Resolve an address string into a net.TCPAddr.
func resolveBindAddr(bindAddr string) (*net.TCPAddr, error) {
	addr, err := net.ResolveTCPAddr("tcp", bindAddr)
	if err == nil {
		return addr, nil
	}
	// Before the fixing of bug #7011, tor doesn't put brackets around IPv6
	// addresses. Split after the last colon, assuming it is a port
	// separator, and try adding the brackets.
	parts := strings.Split(bindAddr, ":")
	if len(parts) <= 2 {
		return nil, err
	}
	bindAddr = "[" + strings.Join(parts[:len(parts)-1], ":") + "]:" + parts[len(parts)-1]
	return net.ResolveTCPAddr("tcp", bindAddr)
}

// Return a new slice, the members of which are those members of addrs having a
// MethodName in methodsNames.
func filterBindAddrs(addrs []PtBindAddr, methodNames []string) []PtBindAddr {
	var result []PtBindAddr

	for _, ba := range addrs {
		for _, methodName := range methodNames {
			if ba.MethodName == methodName {
				result = append(result, ba)
				break
			}
		}
	}

	return result
}

// Return a map from method names to bind addresses. The map is the contents of
// TOR_PT_SERVER_BINDADDR, with keys filtered by TOR_PT_SERVER_TRANSPORTS, and
// further filtered by the methods in methodNames.
func getServerBindAddrs(methodNames []string) []PtBindAddr {
	var result []PtBindAddr

	// Get the list of all requested bindaddrs.
	var serverBindAddr = getenvRequired("TOR_PT_SERVER_BINDADDR")
	for _, spec := range strings.Split(serverBindAddr, ",") {
		var bindAddr PtBindAddr

		parts := strings.SplitN(spec, "-", 2)
		if len(parts) != 2 {
			PtEnvError(fmt.Sprintf("TOR_PT_SERVER_BINDADDR: %q: doesn't contain \"-\"", spec))
		}
		bindAddr.MethodName = parts[0]
		addr, err := resolveBindAddr(parts[1])
		if err != nil {
			PtEnvError(fmt.Sprintf("TOR_PT_SERVER_BINDADDR: %q: %s", spec, err.Error()))
		}
		bindAddr.Addr = addr
		result = append(result, bindAddr)
	}

	// Filter by TOR_PT_SERVER_TRANSPORTS.
	serverTransports := getenvRequired("TOR_PT_SERVER_TRANSPORTS")
	if serverTransports != "*" {
		result = filterBindAddrs(result, strings.Split(serverTransports, ","))
	}

	// Finally filter by what we understand.
	result = filterBindAddrs(result, methodNames)

	return result
}

// This structure is returned by PtServerSetup. It consists of a list of
// PtBindAddrs, along with a single address for the ORPort.
type PtServerInfo struct {
	BindAddrs []PtBindAddr
	OrAddr    *net.TCPAddr
}

// Check the server pluggable transports environments, emitting an error message
// and exiting the program if any error is encountered. Resolves the various
// requested bind addresses and the server ORPort. Returns a PtServerInfo
// struct.
func PtServerSetup(methodNames []string) PtServerInfo {
	var info PtServerInfo
	var err error

	ver := getManagedTransportVer()
	if ver == "" {
		PtVersionError("no-version")
	} else {
		PtLine("VERSION", ver)
	}

	var orPort = getenvRequired("TOR_PT_ORPORT")
	info.OrAddr, err = net.ResolveTCPAddr("tcp", orPort)
	if err != nil {
		PtEnvError(fmt.Sprintf("cannot resolve TOR_PT_ORPORT %q: %s", orPort, err.Error()))
	}

	info.BindAddrs = getServerBindAddrs(methodNames)
	if len(info.BindAddrs) == 0 {
		PtSmethodsDone()
		os.Exit(1)
	}

	return info
}
<|endoftext|>


<commit_before>package models

import (
    "goblin/db"
    _ "github.com/go-sql-driver/mysql"
    "log"
    "time"
)

var DB *db.DB
var err error

func init() {
    DB, err = db.Open("mysql", "root:@tcp(127.0.0.1:3306)/sumi")
    if err != nil {
        log.Fatal(err)
    }
    // Migration
    // DB.Migrate(&Book{}, &Author{})
    // DB.DropTable(&Test{})
    // DB.ModifyColumn(&Test{}, "authorid", "int(11) unsigned  COMMENT '修改后作者ID'")
    // DB.DropColumn(&Test{}, "test1")
    // DB.AddIndex(&Test{}, "idx_authorid_created_at", "authorid", "created_at")
    // DB.AddUniqueIndex(&Test{}, "idx_name", "name")
    // DB.RemoveIndex(&Test{}, "name")

    // CRUD
    author = &Author{Name: "测试4", Created_at: time.Now(), Updated_at: time.Now()}
    DB.Create(author)
}
<commit_msg>test Find()<commit_after>package models

import (
    "goblin/db"
    _ "github.com/go-sql-driver/mysql"
    "log"
    // "time"
)

var DB *db.DB
var err error

func init() {
    DB, err = db.Open("mysql", "root:@tcp(127.0.0.1:3306)/sumi?parseTime=true")
    if err != nil {
        log.Fatal(err)
    }
    // Migration
    // DB.Migrate(&Book{}, &Author{})
    // DB.DropTable(&Test{})
    // DB.ModifyColumn(&Test{}, "authorid", "int(11) unsigned  COMMENT '修改后作者ID'")
    // DB.DropColumn(&Test{}, "test1")
    // DB.AddIndex(&Test{}, "idx_authorid_created_at", "authorid", "created_at")
    // DB.AddUniqueIndex(&Test{}, "idx_name", "name")
    // DB.RemoveIndex(&Test{}, "name")

    // CRUD
    // author := &Author{Name: "测试4", Created_at: time.Now(), Updated_at: time.Now()}
    // DB.Create(author)

    author    := Author{}
    rows, err := DB.First(&author)

    // 希望达成的效果
    // log.Printf("author name is: %v", author.Name)

    // 现在的效果
    if err != nil {
        log.Fatal(err)
    }
    for rows.Next() {
        err := rows.Scan(&author.Id, &author.Name, &author.Created_at, &author.Updated_at)
        if err != nil {
            log.Fatal(err)
        }
    }
    log.Printf("result: %v", author)
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"golang.org/x/net/context"
	"os"
	"strconv"
	"time"
	"github.com/digitalocean/godo"
	"golang.org/x/oauth2"
	"flag"
	"fmt"
)

func main() {

	var token string
	flag.StringVar(&token, "token", "", "[REQUIRED] your digital ocean api token")
	flag.Parse()
	if token == "" {
		flag.Usage()
		os.Exit(1)
	}

	if len(flag.Args()) < 1 {
		fmt.Println("command missed!")
		print(help())
		os.Exit(1)
	}

	command := flag.Arg(0)

	tokenSource := &TokenSource{
		AccessToken: token,
	}
	oauthClient := oauth2.NewClient(context.Background(), tokenSource)
	client := godo.NewClient(oauthClient)

	ctx := context.TODO()

	var err error
	if command == "listSnapshots" {
		err = commandListSnapshots(ctx, client)
	} else if command == "snapshotVolume" {
		volumeName := flag.Arg(1)
		if volumeName == "" {
			fmt.Println("volumeName missed!")
			os.Exit(1)
		}
		region := flag.Arg(2)
		if region == "" {
			fmt.Println("region missed!")
			os.Exit(1)
		}
		var snapshotId string
		snapshotId, err = commandSnapshotVolume(ctx, client, volumeName, region, flag.Arg(3))
		if err == nil {
			fmt.Println("snapshot id: \"", snapshotId, "\" for volume: \""+volumeName+"\"  created:")
		}
	} else if command == "deleteSnapshot" {
		snapshotId := flag.Arg(1)
		if snapshotId == "" {
			fmt.Println("snapshotId missed!")
			os.Exit(1)
		}
		commandDeleteSnapshot(ctx, client, snapshotId)
		if err == nil {
			fmt.Println("snapshot id: \"", snapshotId, "\" deleted")
		}
	} else {
		err = errors.New("unknown command: " + command)
	}

	if err != nil {
		os.Stderr.WriteString("ERROR: ", err.Error())
		os.Exit(1)
	}

}

func commandDeleteSnapshot(ctx context.Context, client *godo.Client, snapshotId string) error {
	_, err := client.Snapshots.Delete(ctx, snapshotId)
	return err
}

func commandSnapshotVolume(ctx context.Context, client *godo.Client, volumeName string, region string, snapshotName string) (string, error) {
	volumeId, err := getVolumeIdByVolumeName(ctx, client, volumeName, region)
	if err != nil {
		return "", err
	}
	if snapshotName == "" {
		timestamp := strconv.Itoa(int(time.Now().Unix()))
		snapshotName = volumeName + "-" + timestamp
	}
	request := &godo.SnapshotCreateRequest{}
	request.VolumeID = volumeId
	request.Name = snapshotName

	snapshot, _, err := client.Storage.CreateSnapshot(ctx, request)
	if err != nil {
		return "", err
	}

	return snapshot.ID, nil
}

func getVolumeIdByVolumeName(ctx context.Context, client *godo.Client, volumeName string, region string) (string, error) {

	options := &godo.ListOptions{}
	options.PerPage = 2
	params := &godo.ListVolumeParams{}
	params.Name = volumeName
	params.Region = region
	params.ListOptions = options

	volumes, _, err := client.Storage.ListVolumes(ctx, params)
	if err != nil {
		return "", err
	}
	if len(volumes) > 1 {
		return "", errors.New(strconv.Itoa(len(volumes)) + " volumes found with name " + volumeName + " please provide --region")
	}
	if len(volumes) == 0 {
		return "", errors.New("volume " + volumeName + " not found")
	}
	volumeId := volumes[0].ID
	return volumeId, nil
}

func commandListSnapshots(ctx context.Context, client *godo.Client) error {
	result, err := getSnapshotList(ctx, client)
	if err != nil {
		return err
	}

	fmt.Println("Snaphots found: ", len(result))
	for _, snapshot := range result {
		fmt.Println(snapshot.Name)
		// godo.Snapshot{ID:"28015723", Name:"git.hdws.de 2017-09-22", ResourceID:"5171268", ResourceType:"droplet", Regions:["fra1"], MinDiskSize:30, SizeGigaBytes:5.89, Created:"2017-09-22T22:39:52Z"}
		fmt.Printf("%v %v %v %v[%v] %vGB(%v)\n", snapshot.ResourceType, snapshot.ResourceID, snapshot.Created, snapshot.Name, snapshot.ID, snapshot.SizeGigaBytes, snapshot.MinDiskSize)
	}
	return nil
}

func help() string {
	var help string
	help += "Commands:\n"
	help += "listSnapshots: list all droplet and volume snapshots\n"
	help += "snapshotVolume VOLUMENAME REGION [SNAPSHOTNAME]: create snapshot of given volume\n"
	help += "deleteSnapshot SNAPSHOTID: delete snapshot\n"
	return help
}

type TokenSource struct {
	AccessToken string
}

func (t *TokenSource) Token() (*oauth2.Token, error) {
	token := &oauth2.Token{
		AccessToken: t.AccessToken,
	}
	return token, nil
}

func getSnapshotList(ctx context.Context, client *godo.Client) ([]godo.Snapshot, error) {

	list := []godo.Snapshot{}

	options := &godo.ListOptions{}
	options.PerPage = 1000
	for {
		snapshots, resp, err := client.Snapshots.List(ctx, options)
		if err != nil {
			return nil, err
		}
		// append our list
		for _, d := range snapshots {
			list = append(list, d)
		}

		// if we are at the last page, break out the for loop
		if resp.Links == nil || resp.Links.IsLastPage() {
			break
		}

		page, err := resp.Links.CurrentPage()
		if err != nil {
			return nil, err
		}

		// set the page we want for the next request
		options.Page = page + 1
	}

	return list, nil
}
<commit_msg>add cleanupSnapshots<commit_after>package main

import (
	"errors"
	"golang.org/x/net/context"
	"os"
	"strconv"
	"time"
	"github.com/digitalocean/godo"
	"golang.org/x/oauth2"
	"flag"
	"fmt"
	"regexp"
)

func main() {

	var token string
	flag.StringVar(&token, "token", "", "[REQUIRED] your digital ocean api token")
	flag.Parse()
	if token == "" {
		flag.Usage()
		os.Exit(1)
	}

	if len(flag.Args()) < 1 {
		fmt.Println("command missed!")
		print(help())
		os.Exit(1)
	}

	command := flag.Arg(0)

	tokenSource := &TokenSource{
		AccessToken: token,
	}
	oauthClient := oauth2.NewClient(context.Background(), tokenSource)
	client := godo.NewClient(oauthClient)

	ctx := context.TODO()

	var err error
	if command == "listSnapshots" {
		err = commandListSnapshots(ctx, client)
	} else if command == "snapshotVolume" {
		volumeName := flag.Arg(1)
		if volumeName == "" {
			fmt.Println("volumeName missed!")
			os.Exit(1)
		}
		region := flag.Arg(2)
		if region == "" {
			fmt.Println("region missed!")
			os.Exit(1)
		}
		var snapshotId string
		snapshotId, err = commandSnapshotVolume(ctx, client, volumeName, region, flag.Arg(3))
		if err == nil {
			fmt.Println("snapshot id: \"", snapshotId, "\" for volume: \""+volumeName+"\"  created:")
		}
	} else if command == "deleteSnapshot" {
		snapshotId := flag.Arg(1)
		if snapshotId == "" {
			fmt.Println("snapshotId missed!")
			os.Exit(1)
		}
		commandDeleteSnapshot(ctx, client, snapshotId)
		if err == nil {
			fmt.Println("snapshot id: \"", snapshotId, "\" deleted")
		}
	} else if command == "cleanupSnapshots" {
		resourceId := flag.Arg(1)
		if resourceId == "" {
			fmt.Println("recourceId missed!")
			os.Exit(1)
		}
		retentionString := flag.Arg(2)
		err = commandCleanupSnapshots(ctx, client, resourceId, retentionString)
		if err == nil {
			//fmt.Println("snapshot id: \"", snapshotId, "\" deleted")
		}
	} else {
		err = errors.New("unknown command: " + command)
	}

	if err != nil {
		os.Stderr.WriteString( "Error: " + err.Error() + "\n")
		os.Exit(1)
	}

}

func commandDeleteSnapshot(ctx context.Context, client *godo.Client, snapshotId string) error {
	_, err := client.Snapshots.Delete(ctx, snapshotId)
	return err
}

func commandSnapshotVolume(ctx context.Context, client *godo.Client, volumeName string, region string, snapshotName string) (string, error) {
	volumeId, err := getVolumeIdByVolumeName(ctx, client, volumeName, region)
	if err != nil {
		return "", err
	}
	if snapshotName == "" {
		timestamp := strconv.Itoa(int(time.Now().Unix()))
		snapshotName = volumeName + "-" + timestamp
	}
	request := &godo.SnapshotCreateRequest{}
	request.VolumeID = volumeId
	request.Name = snapshotName

	snapshot, _, err := client.Storage.CreateSnapshot(ctx, request)
	if err != nil {
		return "", err
	}

	return snapshot.ID, nil
}

func getVolumeIdByVolumeName(ctx context.Context, client *godo.Client, volumeName string, region string) (string, error) {

	options := &godo.ListOptions{}
	options.PerPage = 2
	params := &godo.ListVolumeParams{}
	params.Name = volumeName
	params.Region = region
	params.ListOptions = options

	volumes, _, err := client.Storage.ListVolumes(ctx, params)
	if err != nil {
		return "", err
	}
	if len(volumes) > 1 {
		return "", errors.New(strconv.Itoa(len(volumes)) + " volumes found with name " + volumeName + " please provide --region")
	}
	if len(volumes) == 0 {
		return "", errors.New("volume " + volumeName + " not found")
	}
	volumeId := volumes[0].ID
	return volumeId, nil
}

func commandListSnapshots(ctx context.Context, client *godo.Client) error {
	result, err := getSnapshotList(ctx, client)
	if err != nil {
		return err
	}

	fmt.Println("Snaphots found: ", len(result))
	for _, snapshot := range result {
		fmt.Println(snapshot.Name)
		fmt.Printf("  %v %v %v %v[%v] %vGB(%v)\n", snapshot.ResourceType, snapshot.ResourceID, snapshot.Created, snapshot.Name, snapshot.ID, snapshot.SizeGigaBytes, snapshot.MinDiskSize)
	}
	return nil
}

func commandCleanupSnapshots(ctx context.Context, client *godo.Client, resourceId string, retentionString string) error {

	if retentionString == "" {
		retentionString = "1r"
	}
	// retentionString is simple numeric, so we take it as n-last
	if _, err := strconv.Atoi(retentionString); err == nil {
		retentionString = retentionString+"r"
	}

	fullStringRegexp := regexp.MustCompile("^(\\d+[rhdwmy])+$")
	if( !fullStringRegexp.MatchString(retentionString) ){
		return errors.New("invalid retention parameter")
	}

	allSnapshots, err := getSnapshotList(ctx, client)
	if err != nil {
		return err
	}

	snapshots := make([]godo.Snapshot,0);
	for _, snapshot := range allSnapshots {
		if( snapshot.ResourceID==resourceId ) {
			snapshots = append(snapshots,snapshot)
		}
	}

	if len(snapshots)>0 {
		remainingSnapshotIds := make(map[string]bool)
		retentionRegExp := regexp.MustCompile("(\\d+)([rhdwmy]+)")
		for _, retentionElement := range retentionRegExp.FindAllStringSubmatch(retentionString, -1) {
			retentionType := retentionElement[2]
			retentionCount,_ := strconv.Atoi(retentionElement[1])

			switch retentionType {
			case "r":
				startIndex := len(snapshots) - retentionCount
				// emulate max(startIndex, 0)
				if startIndex<0 {
					startIndex = 0
				}
				endIndex := len(snapshots)
				remainingSnapshots := snapshots[startIndex : endIndex]
				for _, snapshot := range remainingSnapshots {
					remainingSnapshotIds[snapshot.ID] = true
				}
			default:
				return errors.New("invalid retentionType " + retentionType)
			}

		}

		for _, snapshot := range snapshots {
			if _,exists := remainingSnapshotIds[snapshot.ID]; exists == false{
				_, err := client.Snapshots.Delete(ctx, snapshot.ID)
				if err != nil {
					return err
				}

			}
		}
	}

	return nil
}

func help() string {
	var help string
	help += "Commands:\n"
	help += "listSnapshots: list all droplet and volume snapshots\n"
	help += "snapshotVolume VOLUMENAME REGION [SNAPSHOTNAME]: create snapshot of given volume\n"
	help += "deleteSnapshot SNAPSHOTID: delete snapshot\n"
	help += "cleanupSnapshots RESOURCEID [RETENTION]: delete all snapshots, not matching to retention-strategy. \n"
	help += "                 retention-strategy: default 1r, int will be converted to (n)r.\n"
	help += "                                  r: 'r' stands for 'recent'. So '3r' means 'keep the 3 most recent snapshots'.\n"
	return help
}

type TokenSource struct {
	AccessToken string
}

func (t *TokenSource) Token() (*oauth2.Token, error) {
	token := &oauth2.Token{
		AccessToken: t.AccessToken,
	}
	return token, nil
}

func getSnapshotList(ctx context.Context, client *godo.Client) ([]godo.Snapshot, error) {

	list := []godo.Snapshot{}

	options := &godo.ListOptions{}
	options.PerPage = 1000
	for {
		snapshots, resp, err := client.Snapshots.List(ctx, options)
		if err != nil {
			return nil, err
		}
		// append our list
		for _, d := range snapshots {
			list = append(list, d)
		}

		// if we are at the last page, break out the for loop
		if resp.Links == nil || resp.Links.IsLastPage() {
			break
		}

		page, err := resp.Links.CurrentPage()
		if err != nil {
			return nil, err
		}

		// set the page we want for the next request
		options.Page = page + 1
	}

	return list, nil
}
<|endoftext|>


<commit_before>package dnsd

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"sync"
	"time"

	"github.com/HouzuoGuo/laitos/lalog"
	"github.com/HouzuoGuo/laitos/misc"
	"github.com/HouzuoGuo/laitos/tcpoverdns"
	"github.com/HouzuoGuo/laitos/toolbox"
)

const (
	// MaxProxyConnectionLifetime is the maximum lifetime of a transmission
	// control. The transmission controls are unconditionally closed after this
	// duration.
	MaxProxyConnectionLifetime = 30 * time.Minute
)

// ProxyRequest is the data sent by a proxy client to initiate a connection
// toward a proxy destination.
type ProxyRequest struct {
	// Network name of the address (e.g. "tcp"), this takes precedence over the
	// port number when the network name is specified.
	Network string `json:"n"`
	// Port number in the absence of network name.
	Port int `json:"p"`
	// Address is the host IP address or network address (IP:port).
	Address string `json:"a"`
	// AccessTOTP is a time-based OTP that authorises the connection request.
	AccessTOTP string `json:"t"`
}

// ProxyConnection consists of a transmission control paired to a TCP connection
// relayed by the transmission control.
type ProxyConnection struct {
	proxy         *Proxy
	tcpConn       *net.TCPConn
	context       context.Context
	tc            *tcpoverdns.TransmissionControl
	buf           *tcpoverdns.SegmentBuffer
	inputSegments net.Conn
	logger        lalog.Logger
}

// Start piping data back and forth between proxy TCP connection and
// transmission control.
// The function blocks until the underlying TC is closed.
func (conn *ProxyConnection) Start() {
	if conn.proxy.Debug {
		conn.logger.Info("", nil, "starting now")
	}
	conn.buf = tcpoverdns.NewSegmentBuffer(conn.logger, conn.tc.Debug, 0)
	beginTimeNano := time.Now().UnixNano()
	defer func() {
		if conn.proxy.Debug {
			conn.logger.Info("", nil, "closing and lingering")
			conn.tc.DumpState()
		}
		_ = conn.Close()
		misc.TCPOverDNSStats.Trigger(float64(time.Now().UnixNano() - beginTimeNano))
		go func() {
			// Linger a short while before deleting (cease tracking) the
			// connection, as the output segment buffer may still contain
			// useful, final few segments - and crucially it contains the
			// segment with ResetTerminate flag which is expected by the peer
			// transmission control.
			time.Sleep(conn.proxy.Linger)
			conn.proxy.mutex.Lock()
			delete(conn.proxy.connections, conn.tc.ID)
			conn.proxy.mutex.Unlock()
			if conn.proxy.Debug {
				conn.logger.Info("", nil, "closed and removed from proxy")
			}
		}()
	}()
	// Absorb outgoing segments into the outgoing backlog.
	conn.tc.OutputSegmentCallback = conn.buf.Absorb
	// Carry on with the handshake.
	conn.tc.Start(conn.context)
	conn.tc.WaitState(conn.context, tcpoverdns.StateEstablished)
	if conn.proxy.Debug {
		conn.logger.Info("", nil, "TC is established")
	}
	// Pipe data in both directions.
	if conn.tcpConn != nil {
		go func() {
			_, err := io.Copy(conn.tcpConn, conn.tc)
			if conn.proxy.Debug {
				conn.logger.Info(nil, err, "finished piping from TC to TCP connection")
			}
		}()
		_, err := io.Copy(conn.tc, conn.tcpConn)
		if conn.proxy.Debug {
			conn.logger.Info(nil, err, "finished piping from TCP connection to TC")
		}
	}
	// Wait for the transmission control to close.
	conn.tc.CloseAfterDrained()
	conn.tc.WaitState(conn.context, tcpoverdns.StateClosed)
	// The proxy connection lingers for a short while, see defer.
}

// WaitSegment busy-waits until a new segment is available from the output
// segment backlog, and then pops the segment.
func (conn *ProxyConnection) WaitSegment(ctx context.Context) (tcpoverdns.Segment, bool) {
	for {
		popped, exists := conn.buf.Pop()
		if exists {
			return popped, true
		}
		select {
		case <-ctx.Done():
			return tcpoverdns.Segment{}, false
		case <-time.After(tcpoverdns.BusyWaitInterval):
			continue
		}
	}
}

// Close and terminate the proxy TCP connection and its transmission control.
func (conn *ProxyConnection) Close() error {
	if conn.tcpConn != nil {
		_ = conn.tcpConn.Close()
	}
	_ = conn.inputSegments.Close()
	_ = conn.tc.Close()
	return nil
}

// Proxy manages the full life cycle of multiple transmission controls created
// for the purpose of relaying TCP connections.
type Proxy struct {
	// DNSDaemon helps the proxy to identify advertising/malware proxy
	// destinations and to refuse serving their clients.
	DNSDaemon *Daemon `json:"-"`

	// RequestOTPSecret is a TOTP secret for authorising incoming connection
	// requests.
	RequestOTPSecret string `json:"RequestOTPSecret"`
	// Debug enables verbose logging for IO activities.
	Debug bool `json:"Debug"`

	// Linger is a brief period of time for a proxy connection to stay before it
	// is removed from the internal collection of proxy connections. The delay
	// is crucial to allow the final segments of each proxy connection to be
	// received by proxy client - including the segment with ResetTerminate.
	Linger time.Duration `json:"-"`

	// DialTimeout is the timeout used for creating new a proxy TCP connection.
	DialTimeout time.Duration `json:"-"`

	// logger is used to log IO activities when verbose logging is enabled.
	logger lalog.Logger `json:"-"`

	connections map[uint16]*ProxyConnection
	context     context.Context
	cancelFun   func()
	mutex       *sync.Mutex
}

// Start initialises the internal state of the proxy.
func (proxy *Proxy) Start(ctx context.Context) {
	if proxy.DialTimeout == 0 {
		proxy.DialTimeout = 10 * time.Second
	}
	if proxy.Linger == 0 {
		proxy.Linger = 60 * time.Second
	}
	proxy.connections = make(map[uint16]*ProxyConnection)
	proxy.context, proxy.cancelFun = context.WithCancel(ctx)
	proxy.mutex = new(sync.Mutex)
	proxy.logger = lalog.Logger{ComponentName: "TCProxy"}
}

// Receive processes an incoming segment and relay the segment to an existing
// transmission control, or create a new transmission control for the proxy
// destination.
func (proxy *Proxy) Receive(in tcpoverdns.Segment) (tcpoverdns.Segment, bool) {
	proxy.mutex.Lock()
	conn, exists := proxy.connections[in.ID]
	proxy.mutex.Unlock()
	if !exists {
		// Connect to the proxy destination.
		var req ProxyRequest
		proxy.logger.Info(in.ID, nil, "new connection request - seg: %+v, req: %+v", in, req)
		if len(in.Data) < tcpoverdns.InitiatorConfigLen {
			proxy.logger.Warning(in.ID, nil, "received a malformed segment possibly from a stale TC")
			return tcpoverdns.Segment{}, false
		}
		if err := json.Unmarshal(in.Data[tcpoverdns.InitiatorConfigLen:], &req); err != nil {
			proxy.logger.Warning(in.ID, err, "failed to deserialise proxy request")
			return tcpoverdns.Segment{}, false
		}
		if proxy.DNSDaemon != nil && proxy.DNSDaemon.IsInBlacklist(req.Address) {
			proxy.logger.Info(in.ID, nil, "refusing connection to blacklisted destination %q", req.Address)
			return tcpoverdns.Segment{ID: in.ID, Flags: tcpoverdns.FlagReset}, true
		}
		prev, curr, next, err := toolbox.GetTwoFACodes(proxy.RequestOTPSecret)
		if err != nil {
			proxy.logger.Info(in.ID, nil, "failed to calculate TOTP codes")
			return tcpoverdns.Segment{ID: in.ID, Flags: tcpoverdns.FlagReset}, true
		}
		if req.AccessTOTP != prev && req.AccessTOTP != curr && req.AccessTOTP != next {
			proxy.logger.Warning(in.ID, nil, "the request failed OTP check")
			return tcpoverdns.Segment{ID: in.ID, Flags: tcpoverdns.FlagReset}, true
		}
		// Construct the transmission control at proxy's side.
		proxyIn, tcIn := net.Pipe()
		// Connect to the intended destination.
		var dialNet, dialDest string
		if req.Network == "" {
			dialNet = "tcp"
			dialDest = fmt.Sprintf("%s:%d", req.Address, req.Port)
		} else {
			dialNet = req.Network
			dialDest = req.Address
		}
		netConn, err := net.DialTimeout(dialNet, dialDest, proxy.DialTimeout)
		if err != nil {
			// Immediately close the transmission control if the destination is
			// unreachable.
			proxy.logger.Warning(in.ID, err, "failed to connect to proxy destination %s %s", dialNet, dialDest)
			// Proceed with handshake, but there will be no data coming through
			// the transmission control and it will be closed shortly.
		}
		var tcpConn *net.TCPConn
		if netConn != nil {
			tcpConn = netConn.(*net.TCPConn)
			misc.TweakTCPConnection(tcpConn, 30*time.Minute)
		}
		// Track the new proxy connection.
		conn = &ProxyConnection{
			proxy:         proxy,
			tcpConn:       tcpConn,
			context:       proxy.context,
			inputSegments: proxyIn,
			logger: lalog.Logger{
				ComponentName: "ProxyConnection",
				ComponentID: []lalog.LoggerIDField{
					{Key: "TCID", Value: in.ID},
					{Key: "Local", Value: tcpConn.LocalAddr().String()},
					{Key: "Remote", Value: tcpConn.RemoteAddr().String()},
				},
			},
		}
		tc := &tcpoverdns.TransmissionControl{
			Debug:  proxy.Debug,
			LogTag: fmt.Sprintf("ProxyConn(%s->%s)", tcpConn.LocalAddr(), tcpConn.RemoteAddr()),
			ID:     in.ID,
			// This transmission control is a responder during the handshake.
			Initiator:      false,
			InputTransport: tcIn,
			MaxLifetime:    MaxProxyConnectionLifetime,
			// In practice there are occasionally bursts of tens of errors at a
			// time before recovery.
			MaxTransportErrors: 300,
			// The duration of all retransmissions (if all go unacknowledged) is
			// MaxRetransmissions x SlidingWindowWaitDuration.
			MaxRetransmissions: 300,
			// The output transport is not used. Instead, the output segments
			// are kept in a backlog.
			OutputTransport: ioutil.Discard,
			// The segment length and sliding window length are set by the
			// initiator using InitiatorConfig.
		}
		tc.PostConfigCallback = func() {
			// After completing handshake and applying the initiator's desired
			// config, tell the segment buffer the desired max. segment length.
			// The buffer will then be able to merge adjacent short segments.
			conn.buf.SetParameters(tc.MaxSegmentLenExclHeader, tc.Debug)
		}
		conn.tc = tc
		proxy.mutex.Lock()
		proxy.connections[in.ID] = conn
		proxy.mutex.Unlock()
		// Go ahead with handshake and data transfer.
		go conn.Start()
	}
	if _, err := conn.inputSegments.Write(in.Packet()); err != nil {
		_ = conn.Close()
	}
	waitCtx, cancel := context.WithTimeout(proxy.context, conn.tc.LiveTiming.AckDelay*8/7)
	defer cancel()
	begin := time.Now()
	seg, hasSeg := conn.WaitSegment(waitCtx)
	proxy.logger.Info(in.ID, nil, "waited %dms for the outbound segment: %+v", time.Since(begin).Milliseconds(), seg)
	return seg, hasSeg
}

// Close terminates all ongoing transmission controls.
// The function always returns nil.
func (proxy *Proxy) Close() error {
	proxy.mutex.Lock()
	// Terminate all TCs.
	for _, conn := range proxy.connections {
		_ = conn.Close()
	}
	defer proxy.mutex.Unlock()
	// Terminate all proxy connections.
	proxy.cancelFun()
	return nil
}
<commit_msg>tcpoverdns: fix potential nil dereference in proxyconnection<commit_after>package dnsd

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"sync"
	"time"

	"github.com/HouzuoGuo/laitos/lalog"
	"github.com/HouzuoGuo/laitos/misc"
	"github.com/HouzuoGuo/laitos/tcpoverdns"
	"github.com/HouzuoGuo/laitos/toolbox"
)

const (
	// MaxProxyConnectionLifetime is the maximum lifetime of a transmission
	// control. The transmission controls are unconditionally closed after this
	// duration.
	MaxProxyConnectionLifetime = 30 * time.Minute
)

// ProxyRequest is the data sent by a proxy client to initiate a connection
// toward a proxy destination.
type ProxyRequest struct {
	// Network name of the address (e.g. "tcp"), this takes precedence over the
	// port number when the network name is specified.
	Network string `json:"n"`
	// Port number in the absence of network name.
	Port int `json:"p"`
	// Address is the host IP address or network address (IP:port).
	Address string `json:"a"`
	// AccessTOTP is a time-based OTP that authorises the connection request.
	AccessTOTP string `json:"t"`
}

// ProxyConnection consists of a transmission control paired to a TCP connection
// relayed by the transmission control.
type ProxyConnection struct {
	proxy         *Proxy
	tcpConn       *net.TCPConn
	context       context.Context
	tc            *tcpoverdns.TransmissionControl
	buf           *tcpoverdns.SegmentBuffer
	inputSegments net.Conn
	logger        lalog.Logger
}

// Start piping data back and forth between proxy TCP connection and
// transmission control.
// The function blocks until the underlying TC is closed.
func (conn *ProxyConnection) Start() {
	if conn.proxy.Debug {
		conn.logger.Info("", nil, "starting now")
	}
	conn.buf = tcpoverdns.NewSegmentBuffer(conn.logger, conn.tc.Debug, 0)
	beginTimeNano := time.Now().UnixNano()
	defer func() {
		if conn.proxy.Debug {
			conn.logger.Info("", nil, "closing and lingering")
			conn.tc.DumpState()
		}
		_ = conn.Close()
		misc.TCPOverDNSStats.Trigger(float64(time.Now().UnixNano() - beginTimeNano))
		go func() {
			// Linger a short while before deleting (cease tracking) the
			// connection, as the output segment buffer may still contain
			// useful, final few segments - and crucially it contains the
			// segment with ResetTerminate flag which is expected by the peer
			// transmission control.
			time.Sleep(conn.proxy.Linger)
			conn.proxy.mutex.Lock()
			delete(conn.proxy.connections, conn.tc.ID)
			conn.proxy.mutex.Unlock()
			if conn.proxy.Debug {
				conn.logger.Info("", nil, "closed and removed from proxy")
			}
		}()
	}()
	// Absorb outgoing segments into the outgoing backlog.
	conn.tc.OutputSegmentCallback = conn.buf.Absorb
	// Carry on with the handshake.
	conn.tc.Start(conn.context)
	conn.tc.WaitState(conn.context, tcpoverdns.StateEstablished)
	if conn.proxy.Debug {
		conn.logger.Info("", nil, "TC is established")
	}
	// Pipe data in both directions.
	if conn.tcpConn != nil {
		go func() {
			_, err := io.Copy(conn.tcpConn, conn.tc)
			if conn.proxy.Debug {
				conn.logger.Info(nil, err, "finished piping from TC to TCP connection")
			}
		}()
		_, err := io.Copy(conn.tc, conn.tcpConn)
		if conn.proxy.Debug {
			conn.logger.Info(nil, err, "finished piping from TCP connection to TC")
		}
	}
	// Wait for the transmission control to close.
	conn.tc.CloseAfterDrained()
	conn.tc.WaitState(conn.context, tcpoverdns.StateClosed)
	// The proxy connection lingers for a short while, see defer.
}

// WaitSegment busy-waits until a new segment is available from the output
// segment backlog, and then pops the segment.
func (conn *ProxyConnection) WaitSegment(ctx context.Context) (tcpoverdns.Segment, bool) {
	for {
		popped, exists := conn.buf.Pop()
		if exists {
			return popped, true
		}
		select {
		case <-ctx.Done():
			return tcpoverdns.Segment{}, false
		case <-time.After(tcpoverdns.BusyWaitInterval):
			continue
		}
	}
}

// Close and terminate the proxy TCP connection and its transmission control.
func (conn *ProxyConnection) Close() error {
	if conn.tcpConn != nil {
		_ = conn.tcpConn.Close()
	}
	_ = conn.inputSegments.Close()
	_ = conn.tc.Close()
	return nil
}

// Proxy manages the full life cycle of multiple transmission controls created
// for the purpose of relaying TCP connections.
type Proxy struct {
	// DNSDaemon helps the proxy to identify advertising/malware proxy
	// destinations and to refuse serving their clients.
	DNSDaemon *Daemon `json:"-"`

	// RequestOTPSecret is a TOTP secret for authorising incoming connection
	// requests.
	RequestOTPSecret string `json:"RequestOTPSecret"`
	// Debug enables verbose logging for IO activities.
	Debug bool `json:"Debug"`

	// Linger is a brief period of time for a proxy connection to stay before it
	// is removed from the internal collection of proxy connections. The delay
	// is crucial to allow the final segments of each proxy connection to be
	// received by proxy client - including the segment with ResetTerminate.
	Linger time.Duration `json:"-"`

	// DialTimeout is the timeout used for creating new a proxy TCP connection.
	DialTimeout time.Duration `json:"-"`

	// logger is used to log IO activities when verbose logging is enabled.
	logger lalog.Logger `json:"-"`

	connections map[uint16]*ProxyConnection
	context     context.Context
	cancelFun   func()
	mutex       *sync.Mutex
}

// Start initialises the internal state of the proxy.
func (proxy *Proxy) Start(ctx context.Context) {
	if proxy.DialTimeout == 0 {
		proxy.DialTimeout = 10 * time.Second
	}
	if proxy.Linger == 0 {
		proxy.Linger = 60 * time.Second
	}
	proxy.connections = make(map[uint16]*ProxyConnection)
	proxy.context, proxy.cancelFun = context.WithCancel(ctx)
	proxy.mutex = new(sync.Mutex)
	proxy.logger = lalog.Logger{ComponentName: "TCProxy"}
}

// Receive processes an incoming segment and relay the segment to an existing
// transmission control, or create a new transmission control for the proxy
// destination.
func (proxy *Proxy) Receive(in tcpoverdns.Segment) (tcpoverdns.Segment, bool) {
	proxy.mutex.Lock()
	conn, exists := proxy.connections[in.ID]
	proxy.mutex.Unlock()
	if !exists {
		// Connect to the proxy destination.
		var req ProxyRequest
		proxy.logger.Info(in.ID, nil, "new connection request - seg: %+v, req: %+v", in, req)
		if len(in.Data) < tcpoverdns.InitiatorConfigLen {
			proxy.logger.Warning(in.ID, nil, "received a malformed segment possibly from a stale TC")
			return tcpoverdns.Segment{}, false
		}
		if err := json.Unmarshal(in.Data[tcpoverdns.InitiatorConfigLen:], &req); err != nil {
			proxy.logger.Warning(in.ID, err, "failed to deserialise proxy request")
			return tcpoverdns.Segment{}, false
		}
		if proxy.DNSDaemon != nil && proxy.DNSDaemon.IsInBlacklist(req.Address) {
			proxy.logger.Info(in.ID, nil, "refusing connection to blacklisted destination %q", req.Address)
			return tcpoverdns.Segment{ID: in.ID, Flags: tcpoverdns.FlagReset}, true
		}
		prev, curr, next, err := toolbox.GetTwoFACodes(proxy.RequestOTPSecret)
		if err != nil {
			proxy.logger.Info(in.ID, nil, "failed to calculate TOTP codes")
			return tcpoverdns.Segment{ID: in.ID, Flags: tcpoverdns.FlagReset}, true
		}
		if req.AccessTOTP != prev && req.AccessTOTP != curr && req.AccessTOTP != next {
			proxy.logger.Warning(in.ID, nil, "the request failed OTP check")
			return tcpoverdns.Segment{ID: in.ID, Flags: tcpoverdns.FlagReset}, true
		}
		// Construct the transmission control at proxy's side.
		proxyIn, tcIn := net.Pipe()
		// Connect to the intended destination.
		var dialNet, dialDest string
		if req.Network == "" {
			dialNet = "tcp"
			dialDest = fmt.Sprintf("%s:%d", req.Address, req.Port)
		} else {
			dialNet = req.Network
			dialDest = req.Address
		}
		netConn, err := net.DialTimeout(dialNet, dialDest, proxy.DialTimeout)
		if err != nil {
			// Immediately close the transmission control if the destination is
			// unreachable.
			proxy.logger.Warning(in.ID, err, "failed to connect to proxy destination %s %s", dialNet, dialDest)
			// Proceed with handshake, but there will be no data coming through
			// the transmission control and it will be closed shortly.
		}
		var localAddr, remoteAddr string
		var tcpConn *net.TCPConn
		if netConn != nil {
			tcpConn = netConn.(*net.TCPConn)
			misc.TweakTCPConnection(tcpConn, 30*time.Minute)
			localAddr = tcpConn.LocalAddr().String()
			remoteAddr = tcpConn.RemoteAddr().String()
		}
		// Track the new proxy connection.
		conn = &ProxyConnection{
			proxy:         proxy,
			tcpConn:       tcpConn,
			context:       proxy.context,
			inputSegments: proxyIn,
			logger: lalog.Logger{
				ComponentName: "ProxyConnection",
				ComponentID: []lalog.LoggerIDField{
					{Key: "TCID", Value: in.ID},
					{Key: "Local", Value: localAddr},
					{Key: "Remote", Value: remoteAddr},
				},
			},
		}
		tc := &tcpoverdns.TransmissionControl{
			Debug:  proxy.Debug,
			LogTag: fmt.Sprintf("ProxyConn(%s->%s)", localAddr, remoteAddr),
			ID:     in.ID,
			// This transmission control is a responder during the handshake.
			Initiator:      false,
			InputTransport: tcIn,
			MaxLifetime:    MaxProxyConnectionLifetime,
			// In practice there are occasionally bursts of tens of errors at a
			// time before recovery.
			MaxTransportErrors: 300,
			// The duration of all retransmissions (if all go unacknowledged) is
			// MaxRetransmissions x SlidingWindowWaitDuration.
			MaxRetransmissions: 300,
			// The output transport is not used. Instead, the output segments
			// are kept in a backlog.
			OutputTransport: ioutil.Discard,
			// The segment length and sliding window length are set by the
			// initiator using InitiatorConfig.
		}
		tc.PostConfigCallback = func() {
			// After completing handshake and applying the initiator's desired
			// config, tell the segment buffer the desired max. segment length.
			// The buffer will then be able to merge adjacent short segments.
			conn.buf.SetParameters(tc.MaxSegmentLenExclHeader, tc.Debug)
		}
		conn.tc = tc
		proxy.mutex.Lock()
		proxy.connections[in.ID] = conn
		proxy.mutex.Unlock()
		// Go ahead with handshake and data transfer.
		go conn.Start()
	}
	if _, err := conn.inputSegments.Write(in.Packet()); err != nil {
		_ = conn.Close()
	}
	waitCtx, cancel := context.WithTimeout(proxy.context, conn.tc.LiveTiming.AckDelay*8/7)
	defer cancel()
	begin := time.Now()
	seg, hasSeg := conn.WaitSegment(waitCtx)
	proxy.logger.Info(in.ID, nil, "waited %dms for the outbound segment: %+v", time.Since(begin).Milliseconds(), seg)
	return seg, hasSeg
}

// Close terminates all ongoing transmission controls.
// The function always returns nil.
func (proxy *Proxy) Close() error {
	proxy.mutex.Lock()
	// Terminate all TCs.
	for _, conn := range proxy.connections {
		_ = conn.Close()
	}
	defer proxy.mutex.Unlock()
	// Terminate all proxy connections.
	proxy.cancelFun()
	return nil
}
<|endoftext|>


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

package main

import (
	"bytes"
	"encoding/csv"
	"fmt"
	"log"
	"sort"
	"sync"
	"time"
)

type timeSeries struct {
	timestamp  int64
	avgLatency time.Duration
	throughPut int64
}

type TimeSeries []timeSeries

func (t TimeSeries) Swap(i, j int)      { t[i], t[j] = t[j], t[i] }
func (t TimeSeries) Len() int           { return len(t) }
func (t TimeSeries) Less(i, j int) bool { return t[i].timestamp < t[j].timestamp }

type secondPoint struct {
	totalLatency time.Duration
	count        int64
}

type secondPoints struct {
	mu sync.Mutex
	tm map[int64]secondPoint
}

func newSecondPoints() *secondPoints {
	return &secondPoints{tm: make(map[int64]secondPoint)}
}

func (sp *secondPoints) Add(ts time.Time, lat time.Duration) {
	sp.mu.Lock()
	defer sp.mu.Unlock()

	tk := ts.Unix()
	if v, ok := sp.tm[tk]; !ok {
		sp.tm[tk] = secondPoint{totalLatency: lat, count: 1}
	} else {
		v.totalLatency += lat
		v.count += 1
		sp.tm[tk] = v
	}
}

func (sp *secondPoints) getTimeSeries() TimeSeries {
	sp.mu.Lock()
	defer sp.mu.Unlock()

	tslice := make(TimeSeries, len(sp.tm))
	i := 0
	for k, v := range sp.tm {
		tslice[i] = timeSeries{
			timestamp:  k,
			avgLatency: time.Duration(v.totalLatency) / time.Duration(v.count),
			throughPut: v.count,
		}
		i++
		if i%100 == 100 {
			log.Printf("processing timeseries at %d / %d", i, len(sp.tm))
		}
	}
	sort.Sort(tslice)
	return tslice
}

func (ts TimeSeries) String() string {
	buf := new(bytes.Buffer)
	wr := csv.NewWriter(buf)
	if err := wr.Write([]string{"unix_ts", "avg_latency", "throughput"}); err != nil {
		log.Fatal(err)
	}
	rows := [][]string{}
	for i := range ts {
		row := []string{
			fmt.Sprintf("%d", ts[i].timestamp),
			fmt.Sprintf("%s", ts[i].avgLatency),
			fmt.Sprintf("%d", ts[i].throughPut),
		}
		rows = append(rows, row)
	}
	if err := wr.WriteAll(rows); err != nil {
		log.Fatal(err)
	}
	wr.Flush()
	if err := wr.Error(); err != nil {
		log.Fatal(err)
	}
	return fmt.Sprintf("\nSample in one second (unix latency throughput):\n%s", buf.String())
}
<commit_msg>bench: fix logging<commit_after>// Copyright 2016 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"bytes"
	"encoding/csv"
	"fmt"
	"log"
	"sort"
	"sync"
	"time"
)

type timeSeries struct {
	timestamp  int64
	avgLatency time.Duration
	throughPut int64
}

type TimeSeries []timeSeries

func (t TimeSeries) Swap(i, j int)      { t[i], t[j] = t[j], t[i] }
func (t TimeSeries) Len() int           { return len(t) }
func (t TimeSeries) Less(i, j int) bool { return t[i].timestamp < t[j].timestamp }

type secondPoint struct {
	totalLatency time.Duration
	count        int64
}

type secondPoints struct {
	mu sync.Mutex
	tm map[int64]secondPoint
}

func newSecondPoints() *secondPoints {
	return &secondPoints{tm: make(map[int64]secondPoint)}
}

func (sp *secondPoints) Add(ts time.Time, lat time.Duration) {
	sp.mu.Lock()
	defer sp.mu.Unlock()

	tk := ts.Unix()
	if v, ok := sp.tm[tk]; !ok {
		sp.tm[tk] = secondPoint{totalLatency: lat, count: 1}
	} else {
		v.totalLatency += lat
		v.count += 1
		sp.tm[tk] = v
	}
}

func (sp *secondPoints) getTimeSeries() TimeSeries {
	sp.mu.Lock()
	defer sp.mu.Unlock()

	tslice := make(TimeSeries, len(sp.tm))
	i := 0
	log.Printf("getTimeSeries has started on %d results\n", len(sp.tm))
	for k, v := range sp.tm {
		tslice[i] = timeSeries{
			timestamp:  k,
			avgLatency: time.Duration(v.totalLatency) / time.Duration(v.count),
			throughPut: v.count,
		}
		i++
		if i%100 == 0 {
			log.Printf("processing timeseries at %d / %d", i, len(sp.tm))
		}
	}
	sort.Sort(tslice)
	return tslice
}

func (ts TimeSeries) String() string {
	buf := new(bytes.Buffer)
	wr := csv.NewWriter(buf)
	if err := wr.Write([]string{"unix_ts", "avg_latency", "throughput"}); err != nil {
		log.Fatal(err)
	}
	rows := [][]string{}
	for i := range ts {
		row := []string{
			fmt.Sprintf("%d", ts[i].timestamp),
			fmt.Sprintf("%s", ts[i].avgLatency),
			fmt.Sprintf("%d", ts[i].throughPut),
		}
		rows = append(rows, row)
	}
	if err := wr.WriteAll(rows); err != nil {
		log.Fatal(err)
	}
	wr.Flush()
	if err := wr.Error(); err != nil {
		log.Fatal(err)
	}
	return fmt.Sprintf("\nSample in one second (unix latency throughput):\n%s", buf.String())
}
<|endoftext|>


<commit_before>package dal

import (
	"reflect"
	"testing"
	"time"
)

func TestGetPurchases(t *testing.T) {
	db, err := mockupDB()
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()

	purchases, err := GetPurchases(db)
	if err != nil {
		t.Error(err)
	}

	if len(purchases) != 2 {
		t.Error("expected 2 purchases")
	}
	if !reflect.DeepEqual(purchases[0], *mockup1()) {
		t.Error("first retrieved row should match expected", purchases[0], *mockup1())
	}
	if !reflect.DeepEqual(purchases[1], *mockup2()) {
		t.Error("second retrieved row should match expected")
	}
}

func TestAddPurchase(t *testing.T) {
	db, err := mockupDB()
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()

	purchases, err := GetPurchases(db)
	if err != nil {
		t.Error(err)
	}
	if len(purchases) != 2 {
		t.Error("expected 2 purchases")
	}

	newPurchase := Purchase{
		uint64(3), "Rawr", 332.03, time.Now(),
	}
	err = AddPurchase(db, &newPurchase)
	if err != nil {
		t.Error(err)
	}

	purchases, err = GetPurchases(db)
	if err != nil {
		t.Error(err)
	}
	if len(purchases) != 3 {
		t.Error("expected 3 purchases now, AddPurchase failed")
	}
	if newPurchase.Name != purchases[0].Name && newPurchase.Cost != purchases[0].Cost {
		t.Error("new purchase was not saved correctly", newPurchase, purchases[0])
	}
}

func TestGetPurchase(t *testing.T) {
	db, err := mockupDB()
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()

	purchase1, err := GetPurchase(db, 1)
	if err != nil {
		t.Error(err)
	}
	if !reflect.DeepEqual(*purchase1, *mockup1()) {
		t.Error("first retrieved row should match expected")
	}

	purchase2, err := GetPurchase(db, 2)
	if err != nil {
		t.Error(err)
	}
	if !reflect.DeepEqual(*purchase2, *mockup2()) {
		t.Error("second retrieved row should match expected")
	}
}

func TestDeletePurchase(t *testing.T) {
	db, err := mockupDB()
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()

	purchases, err := GetPurchases(db)
	if err != nil {
		t.Error(err)
	}
	if len(purchases) != 2 {
		t.Error("expected 2 purchases initially")
	}

	err = DeletePurchase(db, 1)
	if err != nil {
		t.Error(err)
	}

	purchases, err = GetPurchases(db)
	if err != nil {
		t.Error(err)
	}
	if len(purchases) != 1 {
		t.Error("expected 1 purchase after delete")
	}
	if !reflect.DeepEqual(purchases[0], *mockup2()) {
		t.Error("second retrieved row should match expected")
	}

}
<commit_msg>another test<commit_after>package dal

import (
	"reflect"
	"testing"
	"time"
)

func TestGetPurchases(t *testing.T) {
	db, err := mockupDB()
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()

	purchases, err := GetPurchases(db)
	if err != nil {
		t.Error(err)
	}

	if len(purchases) != 2 {
		t.Error("expected 2 purchases")
	}
	if !reflect.DeepEqual(purchases[0], *mockup1()) {
		t.Error("first retrieved row should match expected", purchases[0], *mockup1())
	}
	if !reflect.DeepEqual(purchases[1], *mockup2()) {
		t.Error("second retrieved row should match expected")
	}
}

func TestAddPurchase(t *testing.T) {
	db, err := mockupDB()
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()

	purchases, err := GetPurchases(db)
	if err != nil {
		t.Error(err)
	}
	if len(purchases) != 2 {
		t.Error("expected 2 purchases")
	}

	newPurchase := Purchase{
		uint64(3), "Rawr", 332.03, time.Now(),
	}
	err = AddPurchase(db, &newPurchase)
	if err != nil {
		t.Error(err)
	}

	purchases, err = GetPurchases(db)
	if err != nil {
		t.Error(err)
	}
	if len(purchases) != 3 {
		t.Error("expected 3 purchases now, AddPurchase failed")
	}
	if newPurchase.Name != purchases[0].Name &&
		newPurchase.Cost != purchases[0].Cost &&
		!newPurchase.TimeBought.Before(time.Now()) {
		t.Error("new purchase was not saved correctly", newPurchase, purchases[0])
	}
}

func TestGetPurchase(t *testing.T) {
	db, err := mockupDB()
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()

	purchase1, err := GetPurchase(db, 1)
	if err != nil {
		t.Error(err)
	}
	if !reflect.DeepEqual(*purchase1, *mockup1()) {
		t.Error("first retrieved row should match expected")
	}

	purchase2, err := GetPurchase(db, 2)
	if err != nil {
		t.Error(err)
	}
	if !reflect.DeepEqual(*purchase2, *mockup2()) {
		t.Error("second retrieved row should match expected")
	}
}

func TestDeletePurchase(t *testing.T) {
	db, err := mockupDB()
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()

	purchases, err := GetPurchases(db)
	if err != nil {
		t.Error(err)
	}
	if len(purchases) != 2 {
		t.Error("expected 2 purchases initially")
	}

	err = DeletePurchase(db, 1)
	if err != nil {
		t.Error(err)
	}

	purchases, err = GetPurchases(db)
	if err != nil {
		t.Error(err)
	}
	if len(purchases) != 1 {
		t.Error("expected 1 purchase after delete")
	}
	if !reflect.DeepEqual(purchases[0], *mockup2()) {
		t.Error("second retrieved row should match expected")
	}

}
<|endoftext|>


<commit_before>package simra

import (
	"io"

	mp3 "github.com/hajimehoshi/go-mp3"
	"github.com/hajimehoshi/oto"
	"golang.org/x/mobile/asset"
)

type audio struct {
	player   *oto.Player
	dec      *mp3.Decoder
	resource asset.File
	isClosed chan bool
}

// Audioer is an interface for treating audio
type Audioer interface {
	Play(resource asset.File, loop bool, doneCallback func(err error)) error
	Stop() error
}

// NewAudio returns new audio instance that implements Audioer interface
func NewAudio() Audioer {
	return &audio{}
}

func (a *audio) Play(resource asset.File, loop bool, doneCallback func(err error)) error {
	dec, err := mp3.NewDecoder(resource)
	if err != nil {
		return err
	}
	a.dec = dec

	player, err := oto.NewPlayer(dec.SampleRate(), 2, 2, 8192)
	if err != nil {
		return err
	}
	a.player = player

	a.resource = resource

	doneChan := make(chan error)
	go func() {
		doneChan <- a.doPlay(player, dec, loop)
	}()

	go func() {
		defer func() {
			err := dec.Close()
			if err != nil {
				LogError(err.Error())
			}
			err = player.Close()
			if err != nil {
				LogError(err.Error())
			}
		}()

		<-doneChan
		doneCallback(err)
	}()

	return nil
}

func (a *audio) doPlay(player *oto.Player, r io.ReadSeeker, loop bool) error {
	var offset, written int64
	var err error
	readByte := (int64)(4096 * 30)
loop:
	for {
	playback:
		for {
			r.Seek(offset, io.SeekStart)
			select {
			case <-a.isClosed:
				readByte = 0
				r.Seek(0, io.SeekEnd)
			default:
				// for non-blocking
			}
			written, err = io.CopyN(player, r, readByte)
			if err != nil || written == 0 {
				// error or EOF
				break playback
			}
			offset += written
		}
		if !loop || err != io.EOF {
			break loop
		}
		offset = 0
	}
	return err
}

func (a *audio) Stop() error {
	a.isClosed <- true
	return nil
}
<commit_msg>buffer size for audio.isClosed channel to avoid blocking.<commit_after>package simra

import (
	"io"

	mp3 "github.com/hajimehoshi/go-mp3"
	"github.com/hajimehoshi/oto"
	"golang.org/x/mobile/asset"
)

type audio struct {
	player   *oto.Player
	dec      *mp3.Decoder
	resource asset.File
	isClosed chan bool
}

// Audioer is an interface for treating audio
type Audioer interface {
	Play(resource asset.File, loop bool, doneCallback func(err error)) error
	Stop() error
}

// NewAudio returns new audio instance that implements Audioer interface
func NewAudio() Audioer {
	return &audio{
		isClosed: make(chan bool, 1),
	}
}

func (a *audio) Play(resource asset.File, loop bool, doneCallback func(err error)) error {
	dec, err := mp3.NewDecoder(resource)
	if err != nil {
		return err
	}
	a.dec = dec

	player, err := oto.NewPlayer(dec.SampleRate(), 2, 2, 8192)
	if err != nil {
		return err
	}
	a.player = player

	a.resource = resource

	doneChan := make(chan error)
	go func() {
		doneChan <- a.doPlay(player, dec, loop)
	}()

	go func() {
		defer func() {
			err := dec.Close()
			if err != nil {
				LogError(err.Error())
			}
			err = player.Close()
			if err != nil {
				LogError(err.Error())
			}
		}()

		<-doneChan
		doneCallback(err)
	}()

	return nil
}

func (a *audio) doPlay(player *oto.Player, r io.ReadSeeker, loop bool) error {
	var offset, written int64
	var err error
	readByte := (int64)(4096 * 30)
loop:
	for {
	playback:
		for {
			r.Seek(offset, io.SeekStart)
			select {
			case <-a.isClosed:
				readByte = 0
				r.Seek(0, io.SeekEnd)
			default:
				// for non-blocking
			}
			written, err = io.CopyN(player, r, readByte)
			if err != nil || written == 0 {
				// error or EOF
				break playback
			}
			offset += written
		}
		if !loop || err != io.EOF {
			break loop
		}
		offset = 0
	}
	return err
}

func (a *audio) Stop() error {
	a.isClosed <- true
	return nil
}
<|endoftext|>


<commit_before>package sample_nop_close

import (
	"bytes"
	"fmt"
	"io"
)

func main() {
	var writer io.Writer = createWriter()
	var writerCloser io.WriteCloser = NopCloser(writer)
	needWriteCloser(writerCloser)
	fmt.Printf("done")
}

func createWriter() io.Writer {
	return bytes.NewBufferString("")
}

func needWriteCloser(writer io.WriteCloser) {
	writer.Close()
}

type nopCloser struct {
	io.Writer
}

func (nopCloser) Close() error { return nil }

func NopCloser(r io.Writer) io.WriteCloser {
	return nopCloser{r}
}
<commit_msg>fix package<commit_after>package main

import (
	"bytes"
	"fmt"
	"io"
)

func main() {
	var writer io.Writer = createWriter()
	var writerCloser io.WriteCloser = NopCloser(writer)
	needWriteCloser(writerCloser)
	fmt.Printf("done")
}

func createWriter() io.Writer {
	return bytes.NewBufferString("")
}

func needWriteCloser(writer io.WriteCloser) {
	writer.Close()
}

type nopCloser struct {
	io.Writer
}

func (nopCloser) Close() error { return nil }

func NopCloser(r io.Writer) io.WriteCloser {
	return nopCloser{r}
}
<|endoftext|>


<commit_before>// Copyright 2015 The Cockroach Authors.
//
// 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.
//
// Author: Marc Berhault (marc@cockroachlabs.com)

package server

import (
	"bytes"
	"fmt"
	"io"
	"net/http"
	"testing"

	"github.com/cockroachdb/cockroach/base"
	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/security"
	"github.com/cockroachdb/cockroach/sql/driver"
	"github.com/cockroachdb/cockroach/testutils"
	"github.com/cockroachdb/cockroach/ts"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/leaktest"
	"github.com/gogo/protobuf/proto"
)

func doHTTPReq(t *testing.T, client *http.Client, method, url string, body proto.Message) (*http.Response, error) {
	var b io.Reader
	if body != nil {
		buf, err := proto.Marshal(body)
		if err != nil {
			t.Fatal(err)
		}
		b = bytes.NewReader(buf)
	}
	req, err := http.NewRequest(method, url, b)
	if err != nil {
		t.Fatalf("%s %s: error building request: %s", method, url, err)
	}
	if b != nil {
		req.Header.Add(util.ContentTypeHeader, util.ProtoContentType)
	}

	return client.Do(req)
}

func sqlForUser(context *base.Context) proto.Message {
	ret := &driver.Request{}
	ret.User = context.User
	return ret
}

// Verify client certificate enforcement and user whitelisting.
func TestSSLEnforcement(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := StartTestServer(t)
	defer s.Stop()

	// HTTPS with client certs for "root".
	rootCertsContext := testutils.NewTestBaseContext(security.RootUser)
	// HTTPS with client certs for "node".
	nodeCertsContext := testutils.NewNodeTestBaseContext()
	// HTTPS with client certs for testuser.
	testCertsContext := testutils.NewTestBaseContext(TestUser)
	// HTTPS without client certs. The user does not matter.
	noCertsContext := testutils.NewTestBaseContext(TestUser)
	noCertsContext.Certs = ""
	// Plain http.
	insecureContext := testutils.NewTestBaseContext(TestUser)
	insecureContext.Insecure = true

	kvGet := &roachpb.GetRequest{}
	kvGet.Key = roachpb.Key("/")

	testCases := []struct {
		method, path string
		body         proto.Message
		ctx          *base.Context
		success      bool // request sent successfully (may be non-200)
		code         int  // http response code
	}{
		// /ui/: basic file server: no auth.
		{"GET", "/index.html", nil, rootCertsContext, true, http.StatusOK},
		{"GET", "/index.html", nil, nodeCertsContext, true, http.StatusOK},
		{"GET", "/index.html", nil, testCertsContext, true, http.StatusOK},
		{"GET", "/index.html", nil, noCertsContext, true, http.StatusOK},
		{"GET", "/index.html", nil, insecureContext, false, -1},

		// /_admin/: server.adminServer: no auth.
		{"GET", healthPath, nil, rootCertsContext, true, http.StatusOK},
		{"GET", healthPath, nil, nodeCertsContext, true, http.StatusOK},
		{"GET", healthPath, nil, testCertsContext, true, http.StatusOK},
		{"GET", healthPath, nil, noCertsContext, true, http.StatusOK},
		{"GET", healthPath, nil, insecureContext, false, -1},

		// /debug/: server.adminServer: no auth.
		{"GET", debugEndpoint + "vars", nil, rootCertsContext, true, http.StatusOK},
		{"GET", debugEndpoint + "vars", nil, nodeCertsContext, true, http.StatusOK},
		{"GET", debugEndpoint + "vars", nil, testCertsContext, true, http.StatusOK},
		{"GET", debugEndpoint + "vars", nil, noCertsContext, true, http.StatusOK},
		{"GET", debugEndpoint + "vars", nil, insecureContext, false, -1},

		// /_status/nodes: server.statusServer: no auth.
		{"GET", statusNodesPrefix, nil, rootCertsContext, true, http.StatusOK},
		{"GET", statusNodesPrefix, nil, nodeCertsContext, true, http.StatusOK},
		{"GET", statusNodesPrefix, nil, testCertsContext, true, http.StatusOK},
		{"GET", statusNodesPrefix, nil, noCertsContext, true, http.StatusOK},
		{"GET", statusNodesPrefix, nil, insecureContext, false, -1},

		// /ts/: ts.Server: no auth.
		{"GET", ts.URLPrefix, nil, rootCertsContext, true, http.StatusNotFound},
		{"GET", ts.URLPrefix, nil, nodeCertsContext, true, http.StatusNotFound},
		{"GET", ts.URLPrefix, nil, testCertsContext, true, http.StatusNotFound},
		{"GET", ts.URLPrefix, nil, noCertsContext, true, http.StatusNotFound},
		{"GET", ts.URLPrefix, nil, insecureContext, false, -1},

		// /sql/: sql.Server. These are proto reqs. The important field is header.User.
		{"POST", driver.Endpoint + driver.Execute.String(), sqlForUser(rootCertsContext),
			rootCertsContext, true, http.StatusOK},
		{"POST", driver.Endpoint + driver.Execute.String(), sqlForUser(nodeCertsContext),
			nodeCertsContext, true, http.StatusOK},
		{"POST", driver.Endpoint + driver.Execute.String(), sqlForUser(testCertsContext),
			testCertsContext, true, http.StatusOK},
		{"POST", driver.Endpoint + driver.Execute.String(), sqlForUser(noCertsContext),
			noCertsContext, true, http.StatusUnauthorized},
		{"POST", driver.Endpoint + driver.Execute.String(), sqlForUser(insecureContext),
			insecureContext, false, -1},
	}

	for tcNum, tc := range testCases {
		client, err := tc.ctx.GetHTTPClient()
		if err != nil {
			t.Fatalf("[%d]: failed to get http client: %v", tcNum, err)
		}
		url := fmt.Sprintf("%s://%s%s", tc.ctx.HTTPRequestScheme(), s.ServingAddr(), tc.path)
		resp, err := doHTTPReq(t, client, tc.method, url, tc.body)
		if (err == nil) != tc.success {
			t.Fatalf("[%d]: expected success=%t, got err=%v", tcNum, tc.success, err)
		}
		if err != nil {
			continue
		}

		defer resp.Body.Close()
		if resp.StatusCode != tc.code {
			t.Errorf("[%d]: expected status code %d, got %d", tcNum, tc.code, resp.StatusCode)
		}
	}
}
<commit_msg>Skip a flaky case in TestSSLEnforcement<commit_after>// Copyright 2015 The Cockroach Authors.
//
// 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.
//
// Author: Marc Berhault (marc@cockroachlabs.com)

package server

import (
	"bytes"
	"fmt"
	"io"
	"net/http"
	"testing"

	"github.com/cockroachdb/cockroach/base"
	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/security"
	"github.com/cockroachdb/cockroach/sql/driver"
	"github.com/cockroachdb/cockroach/testutils"
	"github.com/cockroachdb/cockroach/ts"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/leaktest"
	"github.com/gogo/protobuf/proto"
)

func doHTTPReq(t *testing.T, client *http.Client, method, url string, body proto.Message) (*http.Response, error) {
	var b io.Reader
	if body != nil {
		buf, err := proto.Marshal(body)
		if err != nil {
			t.Fatal(err)
		}
		b = bytes.NewReader(buf)
	}
	req, err := http.NewRequest(method, url, b)
	if err != nil {
		t.Fatalf("%s %s: error building request: %s", method, url, err)
	}
	if b != nil {
		req.Header.Add(util.ContentTypeHeader, util.ProtoContentType)
	}

	return client.Do(req)
}

func sqlForUser(context *base.Context) proto.Message {
	ret := &driver.Request{}
	ret.User = context.User
	return ret
}

// Verify client certificate enforcement and user whitelisting.
func TestSSLEnforcement(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := StartTestServer(t)
	defer s.Stop()

	// HTTPS with client certs for "root".
	rootCertsContext := testutils.NewTestBaseContext(security.RootUser)
	// HTTPS with client certs for "node".
	nodeCertsContext := testutils.NewNodeTestBaseContext()
	// HTTPS with client certs for testuser.
	testCertsContext := testutils.NewTestBaseContext(TestUser)
	// HTTPS without client certs. The user does not matter.
	noCertsContext := testutils.NewTestBaseContext(TestUser)
	noCertsContext.Certs = ""
	// Plain http.
	insecureContext := testutils.NewTestBaseContext(TestUser)
	insecureContext.Insecure = true

	kvGet := &roachpb.GetRequest{}
	kvGet.Key = roachpb.Key("/")

	testCases := []struct {
		method, path string
		body         proto.Message
		ctx          *base.Context
		success      bool // request sent successfully (may be non-200)
		code         int  // http response code
	}{
		// /ui/: basic file server: no auth.
		{"GET", "/index.html", nil, rootCertsContext, true, http.StatusOK},
		{"GET", "/index.html", nil, nodeCertsContext, true, http.StatusOK},
		{"GET", "/index.html", nil, testCertsContext, true, http.StatusOK},
		{"GET", "/index.html", nil, noCertsContext, true, http.StatusOK},
		{"GET", "/index.html", nil, insecureContext, false, -1},

		// /_admin/: server.adminServer: no auth.
		{"GET", healthPath, nil, rootCertsContext, true, http.StatusOK},
		{"GET", healthPath, nil, nodeCertsContext, true, http.StatusOK},
		{"GET", healthPath, nil, testCertsContext, true, http.StatusOK},
		{"GET", healthPath, nil, noCertsContext, true, http.StatusOK},
		{"GET", healthPath, nil, insecureContext, false, -1},

		// /debug/: server.adminServer: no auth.
		{"GET", debugEndpoint + "vars", nil, rootCertsContext, true, http.StatusOK},
		{"GET", debugEndpoint + "vars", nil, nodeCertsContext, true, http.StatusOK},
		{"GET", debugEndpoint + "vars", nil, testCertsContext, true, http.StatusOK},
		{"GET", debugEndpoint + "vars", nil, noCertsContext, true, http.StatusOK},
		{"GET", debugEndpoint + "vars", nil, insecureContext, false, -1},

		// /_status/nodes: server.statusServer: no auth.
		{"GET", statusNodesPrefix, nil, rootCertsContext, true, http.StatusOK},
		{"GET", statusNodesPrefix, nil, nodeCertsContext, true, http.StatusOK},
		{"GET", statusNodesPrefix, nil, testCertsContext, true, http.StatusOK},
		{"GET", statusNodesPrefix, nil, noCertsContext, true, http.StatusOK},
		{"GET", statusNodesPrefix, nil, insecureContext, false, -1},

		// /ts/: ts.Server: no auth.
		{"GET", ts.URLPrefix, nil, rootCertsContext, true, http.StatusNotFound},
		{"GET", ts.URLPrefix, nil, nodeCertsContext, true, http.StatusNotFound},
		{"GET", ts.URLPrefix, nil, testCertsContext, true, http.StatusNotFound},
		{"GET", ts.URLPrefix, nil, noCertsContext, true, http.StatusNotFound},
		{"GET", ts.URLPrefix, nil, insecureContext, false, -1},

		// /sql/: sql.Server. These are proto reqs. The important field is header.User.
		{"POST", driver.Endpoint + driver.Execute.String(), sqlForUser(rootCertsContext),
			rootCertsContext, true, http.StatusOK},
		{"POST", driver.Endpoint + driver.Execute.String(), sqlForUser(nodeCertsContext),
			nodeCertsContext, true, http.StatusOK},
		{"POST", driver.Endpoint + driver.Execute.String(), sqlForUser(testCertsContext),
			testCertsContext, true, http.StatusOK},
		// TODO(bdarnell): https://github.com/cockroachdb/cockroach/issues/4718
		//{"POST", driver.Endpoint + driver.Execute.String(), sqlForUser(noCertsContext),
		//	noCertsContext, true, http.StatusUnauthorized},
		{"POST", driver.Endpoint + driver.Execute.String(), sqlForUser(insecureContext),
			insecureContext, false, -1},
	}

	for tcNum, tc := range testCases {
		client, err := tc.ctx.GetHTTPClient()
		if err != nil {
			t.Fatalf("[%d]: failed to get http client: %v", tcNum, err)
		}
		url := fmt.Sprintf("%s://%s%s", tc.ctx.HTTPRequestScheme(), s.ServingAddr(), tc.path)
		resp, err := doHTTPReq(t, client, tc.method, url, tc.body)
		if (err == nil) != tc.success {
			t.Fatalf("[%d]: expected success=%t, got err=%v", tcNum, tc.success, err)
		}
		if err != nil {
			continue
		}

		defer resp.Body.Close()
		if resp.StatusCode != tc.code {
			t.Errorf("[%d]: expected status code %d, got %d", tcNum, tc.code, resp.StatusCode)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License. See the LICENSE file in the project root for license information.

package arm

import (
	"fmt"

	"github.com/Azure/azure-sdk-for-go/arm/compute"
	"github.com/Azure/packer-azure/packer/builder/azure/common/constants"
	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/packer"
)

type StepCaptureImage struct {
	client  *AzureClient
	capture func(resourceGroupName string, computeName string, parameters *compute.VirtualMachineCaptureParameters) error
	say     func(message string)
	error   func(e error)
}

func NewStepCaptureImage(client *AzureClient, ui packer.Ui) *StepCaptureImage {
	var step = &StepCaptureImage{
		client: client,
		say:    func(message string) { ui.Say(message) },
		error:  func(e error) { ui.Error(e.Error()) },
	}

	step.capture = step.captureImage
	return step
}

func (s *StepCaptureImage) captureImage(resourceGroupName string, computeName string, parameters *compute.VirtualMachineCaptureParameters) error {
	generalizeResponse, err := s.client.Generalize(resourceGroupName, computeName)
	if err != nil {
		return err
	}

	s.client.VirtualMachinesClient.PollAsNeeded(generalizeResponse.Response)

	captureResponse, err := s.client.Capture(resourceGroupName, computeName, *parameters)
	if err != nil {
		return err
	}

	s.client.VirtualMachinesClient.PollAsNeeded(captureResponse.Response.Response)
	return nil
}

func (s *StepCaptureImage) Run(state multistep.StateBag) multistep.StepAction {
	s.say("Capturing image ...")

	var computeName = state.Get(constants.ArmComputeName).(string)
	var resourceGroupName = state.Get(constants.ArmResourceGroupName).(string)
	var parameters = state.Get(constants.ArmVirtualMachineCaptureParameters).(*compute.VirtualMachineCaptureParameters)

	s.say(fmt.Sprintf(" -> ResourceGroupName : '%s'", resourceGroupName))
	s.say(fmt.Sprintf(" -> ComputeName       : '%s'", computeName))

	err := s.capture(resourceGroupName, computeName, parameters)
	if err != nil {
		state.Put(constants.Error, err)
		s.error(err)

		return multistep.ActionHalt
	}

	state.Put(constants.ArmIsValid, true)
	return multistep.ActionContinue
}

func (*StepCaptureImage) Cleanup(multistep.StateBag) {
}
<commit_msg>Respond to review feedback.<commit_after>// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License. See the LICENSE file in the project root for license information.

package arm

import (
	"fmt"

	"github.com/Azure/azure-sdk-for-go/arm/compute"
	"github.com/Azure/packer-azure/packer/builder/azure/common/constants"
	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/packer"
)

type StepCaptureImage struct {
	client  *AzureClient
	capture func(resourceGroupName string, computeName string, parameters *compute.VirtualMachineCaptureParameters) error
	say     func(message string)
	error   func(e error)
}

func NewStepCaptureImage(client *AzureClient, ui packer.Ui) *StepCaptureImage {
	var step = &StepCaptureImage{
		client: client,
		say:    func(message string) { ui.Say(message) },
		error:  func(e error) { ui.Error(e.Error()) },
	}

	step.capture = step.captureImage
	return step
}

func (s *StepCaptureImage) captureImage(resourceGroupName string, computeName string, parameters *compute.VirtualMachineCaptureParameters) error {
	generalizeResponse, err := s.client.Generalize(resourceGroupName, computeName)
	if err != nil {
		return err
	}

	s.client.VirtualMachinesClient.PollAsNeeded(generalizeResponse.Response)

	captureResponse, err := s.client.Capture(resourceGroupName, computeName, *parameters)
	if err != nil {
		return err
	}

	s.client.VirtualMachinesClient.PollAsNeeded(captureResponse.Response.Response)
	return nil
}

func (s *StepCaptureImage) Run(state multistep.StateBag) multistep.StepAction {
	s.say("Capturing image ...")

	var computeName = state.Get(constants.ArmComputeName).(string)
	var resourceGroupName = state.Get(constants.ArmResourceGroupName).(string)
	var parameters = state.Get(constants.ArmVirtualMachineCaptureParameters).(*compute.VirtualMachineCaptureParameters)

	s.say(fmt.Sprintf(" -> ResourceGroupName : '%s'", resourceGroupName))
	s.say(fmt.Sprintf(" -> ComputeName       : '%s'", computeName))

	err := s.capture(resourceGroupName, computeName, parameters)
	if err != nil {
		state.Put(constants.Error, err)
		s.error(err)

		return multistep.ActionHalt
	}

	return multistep.ActionContinue
}

func (*StepCaptureImage) Cleanup(multistep.StateBag) {
}
<|endoftext|>