<commit_before>package sqlgen

import (
	"bytes"
	"errors"
	"fmt"
	"strings"
)

type DatabaseQuery struct {
	action      string
	distinct    bool
	columns     []string
	values      [][]interface{}
	tables      []string
	setParts    []conditionPart
	whereParts  []conditionPart
	havingParts []conditionPart
	groupBys    []string
	orderBys    []string
	limit       string
	offset      string
}

// From "... from ..."
func (q *DatabaseQuery) From(tables ...string) error {
	if q.action != "SELECT" && q.action != "DELETE" {
		return errors.New(q.action + " not support from function")
	}
	if q.action == "DELETE" {
		if len(tables) != 1 || len(q.tables) > 0 {
			return errors.New(q.action + " not support mulit-tables")
		}
	}
	q.tables = append(q.tables, tables...)
	return nil
}

// Tables add tables
func (q *DatabaseQuery) Tables(tables ...string) error {
	if q.action != "SELECT" && q.action != "UPDATE" && len(q.tables) > 0 {
		return errors.New(q.action + " not support mulit-tables")
	}
	q.tables = append(q.tables, tables...)
	return nil
}

// Columns add columns
func (q *DatabaseQuery) Columns(columns ...string) {
	q.columns = append(q.columns, columns...)
}

// ToSQL generate sql string
func (q *DatabaseQuery) ToSQL() (string, []interface{}, error) {
	var sqlbuf bytes.Buffer
	var args []interface{}
	var err error
	sqlbuf.WriteString(q.action + " ")
	switch q.action {
	case "SELECT":
		// select ... from ... where ...
		if q.distinct {
			sqlbuf.WriteString("DISTINCT ")
		}
		columns := strings.Join(q.columns, ", ")
		if len(columns) > 0 {
			if columns[0:2] == ", " {
				columns = columns[2:]
			}
			sqlbuf.WriteString(columns)
		} else {
			return sqlbuf.String(), args, errors.New("no columns")
		}
		if len(q.tables) > 0 {
			sqlbuf.WriteString(" FROM ")
			sqlbuf.WriteString(strings.Join(q.tables, ", "))
		} else {
			return sqlbuf.String(), args, errors.New("no tables")
		}
	case "INSERT":
		// insert into ... values ...
		if len(q.tables) == 1 {
			sqlbuf.WriteString("INTO ")
			sqlbuf.WriteString(q.tables[0])
		} else {
			return sqlbuf.String(), args, errors.New("no tables")
		}
		if len(q.columns) > 0 {
			sqlbuf.WriteString(" ( " + strings.Join(q.columns, ", ") + " ) ")
		}
		if len(q.values) > 0 {
			sqlbuf.WriteString("VALUES ")
			var sqlpart []string
			for _, v := range q.values {
				args = append(args, v...)
				var sqls []string
				for i := 0; i < len(v); i++ {
					sqls = append(sqls, "?")
				}
				sqlpart = append(sqlpart, "("+strings.Join(sqls, ",")+")")
			}
			sqlbuf.WriteString(strings.Join(sqlpart, ","))
		}
	case "UPDATE":
		// update ... set ... where ...
		if len(q.tables) > 0 {
			sqlbuf.WriteString(strings.Join(q.tables, ", "))
		} else {
			return sqlbuf.String(), args, errors.New("no tables")
		}
		sqlbuf.WriteString(" SET ")
		setSql, setArgs := conditionPartsToSQL(q.setParts, ", ")
		sqlbuf.WriteString(setSql)
		args = append(args, setArgs...)
	case "DELETE":
		// delete from ... where
		if len(q.tables) == 1 {
			sqlbuf.WriteString("FROM ")
			sqlbuf.WriteString(q.tables[0])
		} else {
			return sqlbuf.String(), args, errors.New("no tables")
		}
	default:
		return sqlbuf.String(), args, errors.New("not support")
	}
	if len(q.whereParts) > 0 {
		sqlbuf.WriteString(" WHERE ")
		whereSql, whereArgs := conditionPartsToSQL(q.whereParts, " AND ")
		sqlbuf.WriteString(whereSql)
		args = append(args, whereArgs...)
	}
	if len(q.groupBys) > 0 {
		sqlbuf.WriteString(" GROUP BY ")
		sqlbuf.WriteString(strings.Join(q.groupBys, ", "))
	}

	if len(q.havingParts) > 0 {
		sqlbuf.WriteString(" HAVING ")
		havingSql, havingArgs := conditionPartsToSQL(q.havingParts, " AND ")
		sqlbuf.WriteString(havingSql)
		args = append(args, havingArgs...)
	}
	if q.action == "SELECT" {
		if len(q.orderBys) > 0 {
			sqlbuf.WriteString(" ORDER BY ")
			sqlbuf.WriteString(strings.Join(q.orderBys, ", "))
		}
		if len(q.limit) > 0 {
			sqlbuf.WriteString(" LIMIT ")
			sqlbuf.WriteString(q.limit)
		}
		if len(q.offset) > 0 {
			sqlbuf.WriteString(" OFFSET ")
			sqlbuf.WriteString(q.offset)
		}
	}
	return sqlbuf.String(), args, err
}

func PostgresSQLFormat(query string) string {
	parts := strings.Split(query, "?")
	var sqlbuf bytes.Buffer
	for i, v := range parts {
		sqlbuf.WriteString(v)
		if len(parts) != (i + 1) {
			sqlbuf.WriteString(fmt.Sprintf("$%d", i+1))
		}
	}
	return sqlbuf.String()
}
<commit_msg>fix order by<commit_after>package sqlgen

import (
	"bytes"
	"errors"
	"fmt"
	"regexp"
	"strings"
)

type DatabaseQuery struct {
	action      string
	distinct    bool
	columns     []string
	values      [][]interface{}
	tables      []string
	setParts    []conditionPart
	whereParts  []conditionPart
	havingParts []conditionPart
	groupBys    []string
	orderBys    []string
	limit       string
	offset      string
}

// From "... from ..."
func (q *DatabaseQuery) From(tables ...string) error {
	if q.action != "SELECT" && q.action != "DELETE" {
		return errors.New(q.action + " not support from function")
	}
	if q.action == "DELETE" {
		if len(tables) != 1 || len(q.tables) > 0 {
			return errors.New(q.action + " not support mulit-tables")
		}
	}
	q.tables = append(q.tables, tables...)
	return nil
}

// Tables add tables
func (q *DatabaseQuery) Tables(tables ...string) error {
	if q.action != "SELECT" && q.action != "UPDATE" && len(q.tables) > 0 {
		return errors.New(q.action + " not support mulit-tables")
	}
	q.tables = append(q.tables, tables...)
	return nil
}

// Columns add columns
func (q *DatabaseQuery) Columns(columns ...string) {
	q.columns = append(q.columns, columns...)
}

// ToSQL generate sql string
func (q *DatabaseQuery) ToSQL() (string, []interface{}, error) {
	var sqlbuf bytes.Buffer
	var args []interface{}
	var err error
	sqlbuf.WriteString(q.action + " ")
	switch q.action {
	case "SELECT":
		// select ... from ... where ...
		if q.distinct {
			sqlbuf.WriteString("DISTINCT ")
		}
		columns := strings.Join(q.columns, ", ")
		if len(columns) > 0 {
			if columns[0:2] == ", " {
				columns = columns[2:]
			}
			sqlbuf.WriteString(columns)
		} else {
			return sqlbuf.String(), args, errors.New("no columns")
		}
		if len(q.tables) > 0 {
			sqlbuf.WriteString(" FROM ")
			sqlbuf.WriteString(strings.Join(q.tables, ", "))
		} else {
			return sqlbuf.String(), args, errors.New("no tables")
		}
	case "INSERT":
		// insert into ... values ...
		if len(q.tables) == 1 {
			sqlbuf.WriteString("INTO ")
			sqlbuf.WriteString(q.tables[0])
		} else {
			return sqlbuf.String(), args, errors.New("no tables")
		}
		if len(q.columns) > 0 {
			sqlbuf.WriteString(" ( " + strings.Join(q.columns, ", ") + " ) ")
		}
		if len(q.values) > 0 {
			sqlbuf.WriteString("VALUES ")
			var sqlpart []string
			for _, v := range q.values {
				args = append(args, v...)
				var sqls []string
				for i := 0; i < len(v); i++ {
					sqls = append(sqls, "?")
				}
				sqlpart = append(sqlpart, "("+strings.Join(sqls, ",")+")")
			}
			sqlbuf.WriteString(strings.Join(sqlpart, ","))
		}
	case "UPDATE":
		// update ... set ... where ...
		if len(q.tables) > 0 {
			sqlbuf.WriteString(strings.Join(q.tables, ", "))
		} else {
			return sqlbuf.String(), args, errors.New("no tables")
		}
		sqlbuf.WriteString(" SET ")
		setSql, setArgs := conditionPartsToSQL(q.setParts, ", ")
		sqlbuf.WriteString(setSql)
		args = append(args, setArgs...)
	case "DELETE":
		// delete from ... where
		if len(q.tables) == 1 {
			sqlbuf.WriteString("FROM ")
			sqlbuf.WriteString(q.tables[0])
		} else {
			return sqlbuf.String(), args, errors.New("no tables")
		}
	default:
		return sqlbuf.String(), args, errors.New("not support")
	}
	if len(q.whereParts) > 0 {
		sqlbuf.WriteString(" WHERE ")
		whereSql, whereArgs := conditionPartsToSQL(q.whereParts, " AND ")
		sqlbuf.WriteString(whereSql)
		args = append(args, whereArgs...)
	}
	if len(q.groupBys) > 0 {
		sqlbuf.WriteString(" GROUP BY ")
		sqlbuf.WriteString(strings.Join(q.groupBys, ", "))
	}

	if len(q.havingParts) > 0 {
		sqlbuf.WriteString(" HAVING ")
		havingSql, havingArgs := conditionPartsToSQL(q.havingParts, " AND ")
		sqlbuf.WriteString(havingSql)
		args = append(args, havingArgs...)
	}
	if q.action == "SELECT" {
		if len(q.orderBys) > 0 {
			sqlbuf.WriteString(" ORDER BY ")
			valid := regexp.MustCompile("^[A-Za-z0-9_]+$")
			if !valid.MatchString(strings.Join(q.orderBys, "_")) {
				return sqlbuf.String(), args, fmt.Errorf("bad order by args")
			}
			sqlbuf.WriteString(strings.Join(q.orderBys, ", "))
		}
		if len(q.limit) > 0 {
			sqlbuf.WriteString(" LIMIT ")
			sqlbuf.WriteString(q.limit)
		}
		if len(q.offset) > 0 {
			sqlbuf.WriteString(" OFFSET ")
			sqlbuf.WriteString(q.offset)
		}
	}
	return sqlbuf.String(), args, err
}

func PostgresSQLFormat(query string) string {
	parts := strings.Split(query, "?")
	var sqlbuf bytes.Buffer
	for i, v := range parts {
		sqlbuf.WriteString(v)
		if len(parts) != (i + 1) {
			sqlbuf.WriteString(fmt.Sprintf("$%d", i+1))
		}
	}
	return sqlbuf.String()
}
<|endoftext|>


<commit_before>// Copyright 2014-2015 Apptimist, Inc. 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"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"syscall"
	"time"

	"github.com/apptimistco/asn/debug"
	"github.com/apptimistco/asn/debug/mutex"
	"golang.org/x/net/websocket"
)

type nopCloserWriter struct {
	io.Writer
}

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

// NopCloserWriter returns a WriteCloser with a no-op Close method wrapping
// the provided writer.
func NopCloserWriter(w io.Writer) io.WriteCloser {
	return nopCloserWriter{w}
}

type SrvListener struct {
	ln    Listener
	stop  chan struct{}
	done  chan error
	ws    bool
	clean string
}

type Server struct {
	mutex.Mutex
	cmd       *Command
	repos     Repos
	listeners []*SrvListener
	sessions  []*Ses

	listening struct {
		stop chan struct{}
		done chan struct{}
	}
}

func (cmd *Command) Server(args ...string) {
	srv := &Server{
		cmd:       cmd,
		listeners: make([]*SrvListener, 0),
		sessions:  make([]*Ses, 0),
	}
	err := cmd.Cfg.Check(ServerMode)
	defer func() { cmd.Done <- err }()
	if err != nil {
		runtime.Goexit()
	}
	srv.Mutex.Set(cmd.Cfg.Name)
	if err = srv.repos.Set(cmd.Cfg.Dir); err != nil {
		runtime.Goexit()
	}
	srv.repos.Set(cmd.Cfg.Keys)
	defer func() { srv.repos.Reset() }()
	for _, k := range []*UserKeys{
		srv.cmd.Cfg.Keys.Admin,
		srv.cmd.Cfg.Keys.Server,
	} {
		user := srv.repos.users.User(k.Pub.Encr)
		if user == nil {
			user, err = srv.repos.NewUser(k.Pub.Encr)
			if err != nil {
				runtime.Goexit()
			}
			user.cache.Auth().Set(k.Pub.Auth)
			user.cache.Author().Set(k.Pub.Encr)
		}
		user = nil
	}
	if len(args) > 0 {
		// local server command line exec
		var ses Ses
		// FIXME ses.asn.Init()
		// FIXME defer ses.Reset()
		ses.Set(srv)
		ses.Set(&srv.cmd.Cfg)
		ses.Set(&srv.repos)
		ses.Set(srv.ForEachLogin)
		admin := srv.cmd.Cfg.Keys.Admin.Pub.Encr
		ses.Keys.Client.Login = *admin
		ses.asnsrv = true
		ses.user = ses.asn.repos.users.User(admin)
		v := ses.Exec(NewReqString("exec"), cmd.Stdin, args...)
		err, _ = v.(error)
		AckOut(cmd.Stdout, v)
		v = nil
		runtime.Goexit()
	}
	if err = srv.Listen(); err != nil {
		runtime.Goexit()
	}
	cmd.Stdin.Close()
	cmd.Stdout.Close()
	cmd.Stdout = NopCloserWriter(ioutil.Discard)
	cmd.Stderr.Close()
	cmd.Stderr = NopCloserWriter(ioutil.Discard)
	sigch := make(chan os.Signal, 2)
	signal.Notify(sigch, syscall.SIGINT, syscall.SIGTERM, syscall.SIGUSR1)
	defer signal.Stop(sigch)
	for {
		sig := <-sigch
		srv.Diag("caught", sig)
		switch sig {
		case syscall.SIGINT:
			debug.Trace.WriteTo(debug.Log)
			srv.Close()
			srv.Hangup()
			runtime.Goexit()
		case syscall.SIGTERM:
			srv.Close()
		case syscall.SIGUSR1:
			debug.Trace.WriteTo(debug.Log)
		}
	}
}

func (srv *Server) AddListener(l *SrvListener) {
	srv.Lock()
	defer srv.Unlock()
	for _, p := range srv.listeners {
		if p == nil {
			p = l
			return
		}
	}
	srv.listeners = append(srv.listeners, l)
}

func (srv *Server) Close() {
	for i, le := range srv.listeners {
		if le.ws {
			le.ln.Close()
		} else {
			le.stop <- struct{}{}
			<-le.done
		}
		close(le.stop)
		close(le.done)
		le.ln = nil
		srv.listeners[i] = nil
	}
	srv.listeners = nil
}

func (srv *Server) ForEachLogin(f func(*Ses)) {
	srv.Lock()
	defer srv.Unlock()
	for _, ses := range srv.sessions {
		if ses != nil && ses.asn.IsEstablished() {
			f(ses)
		}
	}
}

func (srv *Server) handler(conn net.Conn) {
	var ses Ses
	svc := srv.cmd.Cfg.Keys
	ses.asn.Init()
	ses.Set(&srv.cmd.Cfg)
	ses.Set(&srv.repos)
	ses.Set(srv.ForEachLogin)
	srv.add(&ses)
	ses.asn.Set(conn)
	defer func() {
		r := recover()
		ses.Lock()
		ses.Unlock()
		if r != nil {
			err := r.(error)
			srv.Log(err)
			ses.asn.Diag(debug.Depth(3), err)
		}
		for i := 0; ses.asn.tx.going; i += 1 {
			if i == 0 {
				close(ses.asn.tx.ch)
			}
			time.Sleep(100 * time.Millisecond)
			if i == 3 {
				panic("can't close connection")
			}
		}
		user := srv.repos.users.User(&ses.Keys.Client.Login)
		if user != nil && user.logins > 0 {
			user.logins -= 1
		}
		srv.rm(&ses)
		srv.Log("disconnected", &ses.Keys.Client.Ephemeral)
		ses.Reset()
	}()
	for {
		n, err := conn.Read(ses.Keys.Client.Ephemeral[:])
		if err == nil {
			if n != PubEncrSz {
				panic(Error{"Oops!", "incomplete ephemeral key"})
			}
			break
		}
		if !IsNetTimeout(err) {
			panic(err)
		}
	}
	ses.asn.Set(NewBox(2, srv.cmd.Cfg.Keys.Nonce,
		&ses.Keys.Client.Ephemeral, svc.Server.Pub.Encr,
		svc.Server.Sec.Encr))
	srv.Log("connected", &ses.Keys.Client.Ephemeral)
	for {
		pdu, opened := <-ses.asn.rx.ch
		if !opened {
			runtime.Goexit()
		}
		err := pdu.Open()
		if err != nil {
			pdu.Free()
			panic(err)
		}
		var v Version
		v.ReadFrom(pdu)
		// FIXME to adjust version ... ses.asn.Set(v)
		var id Id
		id.ReadFrom(pdu)
		id.Internal(v)
		ses.asn.time.out = time.Now()
		switch id {
		case AckReqId:
			err = ses.asn.AckerRx(pdu)
		case ExecReqId:
			err = ses.RxExec(pdu)
		case LoginReqId:
			if err = ses.RxLogin(pdu); err != nil {
				// FIXME block subsequent attempts
				// and eventually terminate
			}
		case BlobId:
			if bytes.Equal(ses.Keys.Client.Login.Bytes(),
				svc.Admin.Pub.Encr.Bytes()) ||
				bytes.Equal(ses.Keys.Client.Login.Bytes(),
					svc.Server.Pub.Encr.Bytes()) {
				_, err = ses.asn.repos.Store(&ses, v, nil, pdu)
			} else {
				err = os.ErrPermission
			}
			if err != nil {
				ses.asn.Diag(err)
			}
		default:
			if id >= Nids {
				panic(ErrIncompatible)
			} else {
				panic(ErrUnsupported)
			}
		}
		pdu.Free()
		pdu = nil
		if err != nil {
			panic(err)
		}
	}
}

func (srv *Server) Hangup() {
	srv.Lock()
	for _, ses := range srv.sessions {
		if ses != nil && ses.asn.tx.going {
			close(ses.asn.tx.ch)
		}
	}
	srv.Unlock()
	for {
		active := 0
		srv.Lock()
		for _, ses := range srv.sessions {
			if ses != nil {
				active += 1
			}
		}
		srv.Unlock()
		if active == 0 {
			break
		}
		time.Sleep(100 * time.Millisecond)
	}
	srv.sessions = nil
}

func (srv *Server) Listen() error {
	for _, lurl := range srv.cmd.Cfg.Listen {
		l := &SrvListener{
			stop: make(chan struct{}, 1),
			done: make(chan error, 1),
		}
		switch lurl.Scheme {
		case "tcp":
			addr, err := net.ResolveTCPAddr(lurl.Scheme, lurl.Host)
			if err != nil {
				return err
			}
			l.ln, err = net.ListenTCP(lurl.Scheme, addr)
			if err != nil {
				return err
			}
			srv.AddListener(l)
			srv.Diag("listening on", addr)
			go l.listen(srv)
		case "unix":
			path := UrlPathSearch(lurl.Path)
			os.Remove(path)
			addr, err := net.ResolveUnixAddr(lurl.Scheme, path)
			if err != nil {
				return err
			}
			l.ln, err = net.ListenUnix(lurl.Scheme, addr)
			if err != nil {
				return err
			}
			srv.AddListener(l)
			l.clean = path
			srv.Diag("listening on", addr)
			go l.listen(srv)
		case "ws":
			l.ws = true
			if lurl.Host == "" {
				lurl.Host = ":http"
			}
			addr, err := net.ResolveTCPAddr("tcp", lurl.Host)
			if err != nil {
				return err
			}
			if l.ln, err = net.ListenTCP("tcp", addr); err != nil {
				return err
			}
			srv.AddListener(l)
			f := func(ws *websocket.Conn) {
				srv.handler(ws)
			}
			/*
				FIXME should use a custom handler
				h := func (w http.ResponseWriter, req *http.Request) {
					s := websocket.Server{Handler: websocket.Handler(webHandler)}
					s.ServeHTTP(w, req)
				});
				s := &http.Server{
					Addr:           ":8080",
					Handler:        h,
					ReadTimeout:    10 * time.Second,
					WriteTimeout:   10 * time.Second,
					MaxHeaderBytes: 1 << 20,
				}
				return s.Serve(l)
			*/
			mux := http.NewServeMux()
			mux.Handle(lurl.Path, websocket.Handler(f))
			srv.Diag("listening on", lurl.String())
			go http.Serve(l.ln, mux)
		default:
			err := &Error{lurl.Scheme, "unsupported"}
			srv.Diag(err)
			return err
		}
	}
	return nil
}

func (srv *Server) add(ses *Ses) {
	srv.Lock()
	defer srv.Unlock()
	for i := range srv.sessions {
		if srv.sessions[i] == nil {
			srv.sessions[i] = ses
			return
		}
	}
	srv.sessions = append(srv.sessions, ses)
}

func (srv *Server) rm(ses *Ses) {
	srv.Lock()
	defer srv.Unlock()
	for i := range srv.sessions {
		if srv.sessions[i] == ses {
			srv.sessions[i] = nil
			break
		}
	}
}

func (l *SrvListener) listen(srv *Server) {
	const dl = 100 * time.Millisecond
	for {
		select {
		case <-l.stop:
			err := l.ln.Close()
			if len(l.clean) > 0 {
				os.Remove(l.clean)
			}
			l.done <- err
			return
		default:
			l.ln.SetDeadline(time.Now().Add(dl))
			conn, err := l.ln.Accept()
			if err == nil {
				l.ln.SetDeadline(time.Time{})
				go srv.handler(conn)
			} else if !IsNetOpTimeout(err) {
				srv.Diag("accept", err)
				runtime.Goexit()
			}
		}
	}
}

func IsNetOpTimeout(err error) bool {
	e, ok := err.(*net.OpError)
	return ok && e.Timeout()
}
<commit_msg>start asn goroutines *after* ephemeral read<commit_after>// Copyright 2014-2015 Apptimist, Inc. 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"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"syscall"
	"time"

	"github.com/apptimistco/asn/debug"
	"github.com/apptimistco/asn/debug/mutex"
	"golang.org/x/net/websocket"
)

type nopCloserWriter struct {
	io.Writer
}

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

// NopCloserWriter returns a WriteCloser with a no-op Close method wrapping
// the provided writer.
func NopCloserWriter(w io.Writer) io.WriteCloser {
	return nopCloserWriter{w}
}

type SrvListener struct {
	ln    Listener
	stop  chan struct{}
	done  chan error
	ws    bool
	clean string
}

type Server struct {
	mutex.Mutex
	cmd       *Command
	repos     Repos
	listeners []*SrvListener
	sessions  []*Ses

	listening struct {
		stop chan struct{}
		done chan struct{}
	}
}

func (cmd *Command) Server(args ...string) {
	srv := &Server{
		cmd:       cmd,
		listeners: make([]*SrvListener, 0),
		sessions:  make([]*Ses, 0),
	}
	err := cmd.Cfg.Check(ServerMode)
	defer func() { cmd.Done <- err }()
	if err != nil {
		runtime.Goexit()
	}
	srv.Mutex.Set(cmd.Cfg.Name)
	if err = srv.repos.Set(cmd.Cfg.Dir); err != nil {
		runtime.Goexit()
	}
	srv.repos.Set(cmd.Cfg.Keys)
	defer func() { srv.repos.Reset() }()
	for _, k := range []*UserKeys{
		srv.cmd.Cfg.Keys.Admin,
		srv.cmd.Cfg.Keys.Server,
	} {
		user := srv.repos.users.User(k.Pub.Encr)
		if user == nil {
			user, err = srv.repos.NewUser(k.Pub.Encr)
			if err != nil {
				runtime.Goexit()
			}
			user.cache.Auth().Set(k.Pub.Auth)
			user.cache.Author().Set(k.Pub.Encr)
		}
		user = nil
	}
	if len(args) > 0 {
		// local server command line exec
		var ses Ses
		// FIXME ses.asn.Init()
		// FIXME defer ses.Reset()
		ses.Set(srv)
		ses.Set(&srv.cmd.Cfg)
		ses.Set(&srv.repos)
		ses.Set(srv.ForEachLogin)
		admin := srv.cmd.Cfg.Keys.Admin.Pub.Encr
		ses.Keys.Client.Login = *admin
		ses.asnsrv = true
		ses.user = ses.asn.repos.users.User(admin)
		v := ses.Exec(NewReqString("exec"), cmd.Stdin, args...)
		err, _ = v.(error)
		AckOut(cmd.Stdout, v)
		v = nil
		runtime.Goexit()
	}
	if err = srv.Listen(); err != nil {
		runtime.Goexit()
	}
	cmd.Stdin.Close()
	cmd.Stdout.Close()
	cmd.Stdout = NopCloserWriter(ioutil.Discard)
	cmd.Stderr.Close()
	cmd.Stderr = NopCloserWriter(ioutil.Discard)
	sigch := make(chan os.Signal, 2)
	signal.Notify(sigch, syscall.SIGINT, syscall.SIGTERM, syscall.SIGUSR1)
	defer signal.Stop(sigch)
	for {
		sig := <-sigch
		srv.Diag("caught", sig)
		switch sig {
		case syscall.SIGINT:
			debug.Trace.WriteTo(debug.Log)
			srv.Close()
			srv.Hangup()
			runtime.Goexit()
		case syscall.SIGTERM:
			srv.Close()
		case syscall.SIGUSR1:
			debug.Trace.WriteTo(debug.Log)
		}
	}
}

func (srv *Server) AddListener(l *SrvListener) {
	srv.Lock()
	defer srv.Unlock()
	for _, p := range srv.listeners {
		if p == nil {
			p = l
			return
		}
	}
	srv.listeners = append(srv.listeners, l)
}

func (srv *Server) Close() {
	for i, le := range srv.listeners {
		if le.ws {
			le.ln.Close()
		} else {
			le.stop <- struct{}{}
			<-le.done
		}
		close(le.stop)
		close(le.done)
		le.ln = nil
		srv.listeners[i] = nil
	}
	srv.listeners = nil
}

func (srv *Server) ForEachLogin(f func(*Ses)) {
	srv.Lock()
	defer srv.Unlock()
	for _, ses := range srv.sessions {
		if ses != nil && ses.asn.IsEstablished() {
			f(ses)
		}
	}
}

func (srv *Server) handler(conn net.Conn) {
	var ses Ses
	svc := srv.cmd.Cfg.Keys
	ses.asn.Init()
	ses.Set(&srv.cmd.Cfg)
	ses.Set(&srv.repos)
	ses.Set(srv.ForEachLogin)
	srv.add(&ses)
	defer func() {
		r := recover()
		ses.Lock()
		ses.Unlock()
		if r != nil {
			err := r.(error)
			srv.Log(err)
			ses.asn.Diag(debug.Depth(3), err)
		}
		for i := 0; ses.asn.tx.going; i += 1 {
			if i == 0 {
				close(ses.asn.tx.ch)
			}
			time.Sleep(100 * time.Millisecond)
			if i == 3 {
				panic("can't close connection")
			}
		}
		user := srv.repos.users.User(&ses.Keys.Client.Login)
		if user != nil && user.logins > 0 {
			user.logins -= 1
		}
		srv.rm(&ses)
		srv.Log("disconnected", &ses.Keys.Client.Ephemeral)
		ses.Reset()
	}()
	for {
		n, err := conn.Read(ses.Keys.Client.Ephemeral[:])
		if err == nil {
			if n != PubEncrSz {
				panic(Error{"Oops!", "incomplete ephemeral key"})
			}
			break
		}
		if !IsNetTimeout(err) {
			panic(err)
		}
	}
	ses.asn.Set(NewBox(2, srv.cmd.Cfg.Keys.Nonce,
		&ses.Keys.Client.Ephemeral, svc.Server.Pub.Encr,
		svc.Server.Sec.Encr))
	srv.Log("connected", &ses.Keys.Client.Ephemeral)
	ses.asn.Set(conn)
	for {
		pdu, opened := <-ses.asn.rx.ch
		if !opened {
			runtime.Goexit()
		}
		err := pdu.Open()
		if err != nil {
			pdu.Free()
			panic(err)
		}
		var v Version
		v.ReadFrom(pdu)
		// FIXME to adjust version ... ses.asn.Set(v)
		var id Id
		id.ReadFrom(pdu)
		id.Internal(v)
		ses.asn.time.out = time.Now()
		switch id {
		case AckReqId:
			err = ses.asn.AckerRx(pdu)
		case ExecReqId:
			err = ses.RxExec(pdu)
		case LoginReqId:
			if err = ses.RxLogin(pdu); err != nil {
				// FIXME block subsequent attempts
				// and eventually terminate
			}
		case BlobId:
			if bytes.Equal(ses.Keys.Client.Login.Bytes(),
				svc.Admin.Pub.Encr.Bytes()) ||
				bytes.Equal(ses.Keys.Client.Login.Bytes(),
					svc.Server.Pub.Encr.Bytes()) {
				_, err = ses.asn.repos.Store(&ses, v, nil, pdu)
			} else {
				err = os.ErrPermission
			}
			if err != nil {
				ses.asn.Diag(err)
			}
		default:
			if id >= Nids {
				panic(ErrIncompatible)
			} else {
				panic(ErrUnsupported)
			}
		}
		pdu.Free()
		pdu = nil
		if err != nil {
			panic(err)
		}
	}
}

func (srv *Server) Hangup() {
	srv.Lock()
	for _, ses := range srv.sessions {
		if ses != nil && ses.asn.tx.going {
			close(ses.asn.tx.ch)
		}
	}
	srv.Unlock()
	for {
		active := 0
		srv.Lock()
		for _, ses := range srv.sessions {
			if ses != nil {
				active += 1
			}
		}
		srv.Unlock()
		if active == 0 {
			break
		}
		time.Sleep(100 * time.Millisecond)
	}
	srv.sessions = nil
}

func (srv *Server) Listen() error {
	for _, lurl := range srv.cmd.Cfg.Listen {
		l := &SrvListener{
			stop: make(chan struct{}, 1),
			done: make(chan error, 1),
		}
		switch lurl.Scheme {
		case "tcp":
			addr, err := net.ResolveTCPAddr(lurl.Scheme, lurl.Host)
			if err != nil {
				return err
			}
			l.ln, err = net.ListenTCP(lurl.Scheme, addr)
			if err != nil {
				return err
			}
			srv.AddListener(l)
			srv.Diag("listening on", addr)
			go l.listen(srv)
		case "unix":
			path := UrlPathSearch(lurl.Path)
			os.Remove(path)
			addr, err := net.ResolveUnixAddr(lurl.Scheme, path)
			if err != nil {
				return err
			}
			l.ln, err = net.ListenUnix(lurl.Scheme, addr)
			if err != nil {
				return err
			}
			srv.AddListener(l)
			l.clean = path
			srv.Diag("listening on", addr)
			go l.listen(srv)
		case "ws":
			l.ws = true
			if lurl.Host == "" {
				lurl.Host = ":http"
			}
			addr, err := net.ResolveTCPAddr("tcp", lurl.Host)
			if err != nil {
				return err
			}
			if l.ln, err = net.ListenTCP("tcp", addr); err != nil {
				return err
			}
			srv.AddListener(l)
			f := func(ws *websocket.Conn) {
				srv.handler(ws)
			}
			/*
				FIXME should use a custom handler
				h := func (w http.ResponseWriter, req *http.Request) {
					s := websocket.Server{Handler: websocket.Handler(webHandler)}
					s.ServeHTTP(w, req)
				});
				s := &http.Server{
					Addr:           ":8080",
					Handler:        h,
					ReadTimeout:    10 * time.Second,
					WriteTimeout:   10 * time.Second,
					MaxHeaderBytes: 1 << 20,
				}
				return s.Serve(l)
			*/
			mux := http.NewServeMux()
			mux.Handle(lurl.Path, websocket.Handler(f))
			srv.Diag("listening on", lurl.String())
			go http.Serve(l.ln, mux)
		default:
			err := &Error{lurl.Scheme, "unsupported"}
			srv.Diag(err)
			return err
		}
	}
	return nil
}

func (srv *Server) add(ses *Ses) {
	srv.Lock()
	defer srv.Unlock()
	for i := range srv.sessions {
		if srv.sessions[i] == nil {
			srv.sessions[i] = ses
			return
		}
	}
	srv.sessions = append(srv.sessions, ses)
}

func (srv *Server) rm(ses *Ses) {
	srv.Lock()
	defer srv.Unlock()
	for i := range srv.sessions {
		if srv.sessions[i] == ses {
			srv.sessions[i] = nil
			break
		}
	}
}

func (l *SrvListener) listen(srv *Server) {
	const dl = 100 * time.Millisecond
	for {
		select {
		case <-l.stop:
			err := l.ln.Close()
			if len(l.clean) > 0 {
				os.Remove(l.clean)
			}
			l.done <- err
			return
		default:
			l.ln.SetDeadline(time.Now().Add(dl))
			conn, err := l.ln.Accept()
			if err == nil {
				l.ln.SetDeadline(time.Time{})
				go srv.handler(conn)
			} else if !IsNetOpTimeout(err) {
				srv.Diag("accept", err)
				runtime.Goexit()
			}
		}
	}
}

func IsNetOpTimeout(err error) bool {
	e, ok := err.(*net.OpError)
	return ok && e.Timeout()
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"io"
	"log"
	"net"
	"os"

	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
)

type SSHClient struct { // SSHSession ...?
	User       string
	Host       string
	Agent      net.Conn
	Conn       *ssh.Client
	Session    *ssh.Session // embed this....?
	StdinPipe  io.WriteCloser
	StdoutPipe io.Reader
}

type Command struct {
	Input  string
	Output io.Reader
}

func (s *SSHClient) Connect(host string) error {
	var signers []ssh.Signer

	// TODO: add the keys from ~/ssh/config ..
	// Look for IdentityFiles .. etc...

	// If there's a running SSH Agent, use its Private keys
	sock, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
	if err == nil {
		agent := agent.NewClient(sock)
		agentSigners, err := agent.Signers()
		if err == nil && len(agentSigners) > 0 {
			log.Println("PublicKey:", agentSigners[0].PublicKey())
			signers = append(signers, agentSigners...)
		}
	}

	s.User = "ubuntu"
	config := &ssh.ClientConfig{
		User: s.User,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(signers...),
		},
	}

	conn, err := ssh.Dial("tcp", host, config)
	if err != nil {
		log.Println("naw..")
		return err
	}
	// defer conn.Close()

	sess, err := conn.NewSession()
	if err != nil {
		log.Println("naw..2")
		return err
	}
	s.Session = sess

	// TODO: test env variables this way....
	// probably wont work... so we need a pty

	// modes := ssh.TerminalModes{
	// 	ssh.ECHO:          0,     // disable echoing
	// 	ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
	// 	ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
	// }
	// // Request pseudo terminal
	// if err := sess.RequestPty("xterm", 80, 40, modes); err != nil {
	// 	log.Fatalf("request for pseudo terminal failed: %s", err)
	// }

	// sess.Setenv("HI", "SUP")

	var b bytes.Buffer
	sess.Stdout = &b
	sess.Stderr = &b
	// err = sess.Start("HI=123 ls -la ; echo $X ; echo sup ; HI=wooooo echo $HI")
	err = sess.Start("ls -la")
	if err != nil {
		return err
	}

	log.Println(b.String())

	sess.Wait()

	return nil
}

func (s *SSHClient) Wait() {
	s.Session.Wait()
}

func (s *SSHClient) Read(b []byte) (int, error) {
	return s.StdoutPipe.Read(b)
}

func (s *SSHClient) Write(b []byte) (int, error) {
	return s.StdinPipe.Write(b)
}
<commit_msg>Add ssh.Connect() host parsing & error handling<commit_after>package main

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"net"
	"net/url"
	"os"
	"strings"

	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
)

type SSHClient struct { // SSHSession ...?
	User       string
	Host       string
	Agent      net.Conn
	Conn       *ssh.Client
	Session    *ssh.Session // embed this....?
	StdinPipe  io.WriteCloser
	StdoutPipe io.Reader
	Env        map[string]string
}

type Command struct {
	Input  string
	Output io.Reader
}

type ErrConnect struct {
	Host   string
	Reason string
}

func (e ErrConnect) Error() string {
	return fmt.Sprintf(`Connect("%v"): %v`, e.Host, e.Reason)
}

// Connect connects SSH client to a specified host.
// Expects host in the format "[ssh://]host[:port]", returns error otherwise.
func (c *SSHClient) Connect(host string) error {
	u, err := url.Parse(host)
	if err != nil {
		return ErrConnect{host, err.Error()}
	}

	// net/url parses "example.com" or "localhost" as relative Path instead of Host.
	// If there's no slash in the Path, we're good to swap Host and Path.
	if u.Host == "" && strings.Index(u.Path, "/") == -1 {
		u.Host, u.Path = u.Path, u.Host
	}
	// Error out on unexpected path
	if u.Path != "" {
		return ErrConnect{host, fmt.Sprintf(`Expected empty path, got "%v"`, u.Path)}
	}
	// Add default port, if not set
	if strings.Index(u.Host, ":") == -1 {
		u.Host = u.Host + ":22"
	}
	// Error out on unknown scheme
	if u.Scheme != "" && u.Scheme != "ssh" {
		return ErrConnect{host, fmt.Sprintf(`Expected "ssh://" (or empty) scheme, got "%v"`, u.Scheme)}
	}
	// Add default user, if not set
	if u.User != nil && u.User.String() != "" {
		c.User = u.User.String()
	} else {
		c.User = os.Getenv("USER")
	}

	// TODO: add the keys from ~/ssh/config ..
	// Look for IdentityFiles .. etc...

	var signers []ssh.Signer

	// If there's a running SSH Agent, use its Private keys
	sock, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
	if err == nil {
		agent := agent.NewClient(sock)
		agentSigners, err := agent.Signers()
		if err == nil && len(agentSigners) > 0 {
			log.Println("PublicKey:", agentSigners[0].PublicKey())
			signers = append(signers, agentSigners...)
		}
	}

	config := &ssh.ClientConfig{
		User: c.User,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(signers...),
		},
	}

	conn, err := ssh.Dial("tcp", u.Host, config)
	if err != nil {
		return ErrConnect{host, err.Error()}
	}

	sess, err := conn.NewSession()
	if err != nil {
		return ErrConnect{host, err.Error()}
	}
	c.Session = sess

	// TODO: test env variables this way....
	// probably wont work... so we need a pty

	// modes := ssh.TerminalModes{
	// 	ssh.ECHO:          0,     // disable echoing
	// 	ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
	// 	ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
	// }
	// // Request pseudo terminal
	// if err := sess.RequestPty("xterm", 80, 40, modes); err != nil {
	// 	log.Fatalf("request for pseudo terminal failed: %s", err)
	// }

	// sess.Setenv("HI", "SUP")

	var b bytes.Buffer
	sess.Stdout = &b
	sess.Stderr = &b
	// err = sess.Start("HI=123 ls -la ; echo $X ; echo sup ; HI=wooooo echo $HI")
	err = sess.Start("ls -la")
	if err != nil {
		return ErrConnect{host, err.Error()}
	}

	log.Println(b.String())

	sess.Wait()

	return nil
}

func (s *SSHClient) Wait() {
	s.Session.Wait()
}

func (s *SSHClient) Read(b []byte) (int, error) {
	return s.StdoutPipe.Read(b)
}

func (s *SSHClient) Write(b []byte) (int, error) {
	return s.StdinPipe.Write(b)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"path/filepath"
	"strings"

	"github.com/bogem/id3v2"
	id3 "github.com/mikkyang/id3-go"
	"github.com/mkideal/cli"
)

type tagT struct {
	cli.Helper
	Auto   bool   `cli:"auto" usage:"short and long format flags both are supported"`
	Artist string `cli:"artist" usage:"Artist name"`
	Album  string `cli:"album" usage:"Album name"`
	Title  string `cli:"title" usage:"Song name"`
	File   string `cli:"file" usage:"A specific file to tag"`
}

var tag = &cli.Command{
	Name: "tag",
	Desc: "Set the id3 tags for mp3 files",
	Argv: func() interface{} { return new(tagT) },
	Fn: func(ctx *cli.Context) error {
		argv := ctx.Argv().(*tagT)

		ctx.String("auto=%v, artist=%s, album=%s, title=%s, file=%s\n", argv.Auto, argv.Artist, argv.Album, argv.Title, argv.File)

		if argv.Auto {
			autoTag()
			return nil
		}

		if checkNotEmpty(argv.File) {
			return saveData(argv.File, argv.Artist, argv.Album, argv.Title)
		}

		files, _ := filepath.Glob("*.mp3")
		for i := 0; i < len(files); i++ {

			fmt.Println("\n************************************")
			fmt.Println("File: " + files[i] + ":")

			err := saveData(files[i], argv.Artist, argv.Album, argv.Title)
			if err != nil {
				log.Println("Error!")
			}

		}

		return nil
	},
}

// Use filename to derive the artist and title
func autoTag() {
	files, _ := filepath.Glob("*.mp3")
	for i := 0; i < len(files); i++ {

		fmt.Println("\n************************************")
		fmt.Println("File: " + files[i] + ":")

		if strings.Index(files[i], " - ") == -1 {
			fmt.Println("Unable to determine track name/artist, skipping")
			continue
		}

		info := strings.Replace(files[i], ".mp3", "", 1)
		trackInfo := strings.Split(info, " - ")

		fmt.Println("  Artist: " + trackInfo[0])
		fmt.Println("  Title: " + trackInfo[1])

		err := saveData(files[i], trackInfo[0], "", trackInfo[1])
		if err != nil {
			log.Println("Error!")
		}
	}
}

func saveData(filename string, artist string, album string, title string) error {
	tag, err := id3v2.Open(filename)
	if err != nil {
		mp3File, terr := id3.Open(filename)
		if terr != nil {
			log.Fatal("Error while opening mp3 file: ", terr)
		}
		defer mp3File.Close()

		if checkNotEmpty(artist) {
			mp3File.SetArtist(artist)
		}

		if checkNotEmpty(album) {
			mp3File.SetAlbum(album)
		}

		if checkNotEmpty(title) {
			mp3File.SetTitle(title)
		}

	} else {
		defer tag.Close()

		if checkNotEmpty(artist) {
			tag.SetArtist(artist)
		}

		if checkNotEmpty(album) {
			tag.SetAlbum(album)
		}

		if checkNotEmpty(title) {
			tag.SetTitle(title)
		}

		saveerr := tag.Save()
		if saveerr != nil {
			log.Fatal("error while saving mp3 file ", saveerr)
		}
	}

	return nil
}
<commit_msg>comment out debug line in tag command<commit_after>package main

import (
	"fmt"
	"log"
	"path/filepath"
	"strings"

	"github.com/bogem/id3v2"
	id3 "github.com/mikkyang/id3-go"
	"github.com/mkideal/cli"
)

type tagT struct {
	cli.Helper
	Auto   bool   `cli:"auto" usage:"short and long format flags both are supported"`
	Artist string `cli:"artist" usage:"Artist name"`
	Album  string `cli:"album" usage:"Album name"`
	Title  string `cli:"title" usage:"Song name"`
	File   string `cli:"file" usage:"A specific file to tag"`
}

var tag = &cli.Command{
	Name: "tag",
	Desc: "Set the id3 tags for mp3 files",
	Argv: func() interface{} { return new(tagT) },
	Fn: func(ctx *cli.Context) error {
		argv := ctx.Argv().(*tagT)

		// ctx.String("auto=%v, artist=%s, album=%s, title=%s, file=%s\n", argv.Auto, argv.Artist, argv.Album, argv.Title, argv.File)

		if argv.Auto {
			autoTag()
			return nil
		}

		if checkNotEmpty(argv.File) {
			return saveData(argv.File, argv.Artist, argv.Album, argv.Title)
		}

		files, _ := filepath.Glob("*.mp3")
		for i := 0; i < len(files); i++ {

			fmt.Println("\n************************************")
			fmt.Println("File: " + files[i] + ":")

			err := saveData(files[i], argv.Artist, argv.Album, argv.Title)
			if err != nil {
				log.Println("Error!")
			}

		}

		return nil
	},
}

// Use filename to derive the artist and title
func autoTag() {
	files, _ := filepath.Glob("*.mp3")
	for i := 0; i < len(files); i++ {

		fmt.Println("\n************************************")
		fmt.Println("File: " + files[i] + ":")

		if strings.Index(files[i], " - ") == -1 {
			fmt.Println("Unable to determine track name/artist, skipping")
			continue
		}

		info := strings.Replace(files[i], ".mp3", "", 1)
		trackInfo := strings.Split(info, " - ")

		fmt.Println("  Artist: " + trackInfo[0])
		fmt.Println("  Title: " + trackInfo[1])

		err := saveData(files[i], trackInfo[0], "", trackInfo[1])
		if err != nil {
			log.Println("Error!")
		}
	}
}

func saveData(filename string, artist string, album string, title string) error {
	tag, err := id3v2.Open(filename)
	if err != nil {
		mp3File, terr := id3.Open(filename)
		if terr != nil {
			log.Fatal("Error while opening mp3 file: ", terr)
		}
		defer mp3File.Close()

		if checkNotEmpty(artist) {
			mp3File.SetArtist(artist)
		}

		if checkNotEmpty(album) {
			mp3File.SetAlbum(album)
		}

		if checkNotEmpty(title) {
			mp3File.SetTitle(title)
		}

	} else {
		defer tag.Close()

		if checkNotEmpty(artist) {
			tag.SetArtist(artist)
		}

		if checkNotEmpty(album) {
			tag.SetAlbum(album)
		}

		if checkNotEmpty(title) {
			tag.SetTitle(title)
		}

		saveerr := tag.Save()
		if saveerr != nil {
			log.Fatal("error while saving mp3 file ", saveerr)
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package carchivum

// tar implements the tape archive format.

import (
	"archive/tar"
	"compress/bzip2"
	"compress/gzip"
	"compress/lzw"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"time"

	"github.com/MichaelTJones/walk"
	"github.com/pierrec/lz4"
)

// Tar is a struct for a tar, tape archive.
type Tar struct {
	Car
	*tar.Writer
	Format
	sources []string
}

// NewTar returns an initialized Tar struct ready for use.
func NewTar(n string) *Tar {
	return &Tar{Car: Car{Name: n, t0: time.Now()}, Format: defaultFormat, sources: []string{}}
}

// Create creates a tarfile from the passed src('s) and saves it to the dst.
func (t *Tar) Create(src ...string) (cnt int, err error) {
	// If there isn't a destination, return err
	if t.Name == "" {
		err = fmt.Errorf("destination required to create a tar archive")
		log.Print(err)
		return 0, err
	}
	// If there aren't any sources, return err
	if len(src) == 0 {
		err = fmt.Errorf("a source is required to create a tar archive")
		log.Print(err)
		return 0, err
	}
	t.sources = src
	// See if we can create the destination file before processing
	tball, err := os.OpenFile(t.Name, os.O_RDWR|os.O_CREATE, 0744)
	if err != nil {
		log.Print(err)
		return 0, err
	}
	defer func() {
		cerr := tball.Close()
		if cerr != nil && err == nil {
			log.Print(cerr)
			err = cerr
		}
	}()
	switch t.Format {
	case GzipFmt:
		err = t.CreateGzip(tball)
		if err != nil {
			log.Print(err)
			return 0, err
		}
	case Bzip2Fmt:
		err = fmt.Errorf("Bzip2 compression is not supported")
		return 0, err
	case LZWFmt:
		err = t.CreateZ(tball)
		if err != nil {
			log.Print(err)
			return 0, err
		}
	case LZ4Fmt:
		err = t.CreateLZ4(tball)
		if err != nil {
			log.Print(err)
			return 0, err
		}
	default:
		err = fmt.Errorf("Unsupported compression format: %s", t.Format.String())
		return 0, err
	}
	if t.DeleteArchived {
		err := t.removeFiles()
		if err != nil {
			err = fmt.Errorf("an error was encountered while deleting the archived files; some files may not be deleted: %q", err)
			log.Print(err)
			return 0, err
		}
	}
	t.setDelta()
	return int(t.Car.files), nil
}

func (t *Tar) removeFiles() error {
	for _, file := range t.deleteList {
		err := os.Remove(file)
		if err != nil {
			log.Print(err)
			return err
		}
	}
	return nil
}

// CreateGzip creates a GZip using the passed writer.
func (t *Tar) CreateGzip(w io.Writer) (err error) {
	zw := gzip.NewWriter(w)
	// Close the file with error handling
	defer func() {
		cerr := zw.Close()
		if cerr != nil && err == nil {
			log.Print(cerr)
			err = cerr
		}
	}()
	err = t.writeTar(zw)
	return err
}

// CreateLZW compresses using LZW and LSB order using the passed writer.
// TODO: address order so that it doesn't necessarily default to LSB
func (t *Tar) CreateZ(w io.Writer) (err error) {
	zw := lzw.NewWriter(w, lzw.LSB, 8)
	// Close the file with error handling
	defer func() {
		cerr := zw.Close()
		if cerr != nil && err == nil {
			log.Print(cerr)
			err = cerr
		}
	}()
	err = t.writeTar(zw)
	return err
}

// CreateLZ4 creates a LZ4 compressed tarball using the passed writer.
func (t *Tar) CreateLZ4(w io.Writer) (err error) {
	lzW := lz4.NewWriter(w)
	err = t.writeTar(lzW)
	lzW.Close()
	return err
}

func (t *Tar) writeTar(w io.Writer) (err error) {
	t.Writer = tar.NewWriter(w)
	defer func() {
		cerr := t.Writer.Close()
		if cerr != nil && err == nil {
			log.Print(cerr)
			err = cerr
		}
	}()
	t.FileCh = make(chan *os.File)
	wait, err := t.Write()
	if err != nil {
		log.Print(err)
		return err
	}
	var fullPath string
	visitor := func(p string, fi os.FileInfo, err error) error {
		return t.AddFile(fullPath, p, fi, err)
	}
	for _, source := range t.sources {
		fullPath, err = filepath.Abs(source)
		if err != nil {
			log.Print(err)
			return err
		}
		err = walk.Walk(fullPath, visitor)
		if err != nil {
			log.Print(err)
			return err
		}
	}
	close(t.FileCh)
	wait.Wait()
	return err
}

// Write adds the files received on the channel to the tarball.
func (t *Tar) Write() (*sync.WaitGroup, error) {
	var wg sync.WaitGroup
	wg.Add(1)
	go func() error {
		defer wg.Done()
		for f := range t.FileCh {
			info, err := f.Stat()
			if err != nil {
				log.Print(err)
				return err
			}
			if info.IsDir() {
				continue
			}
			header, err := tar.FileInfoHeader(info, "")
			if err != nil {
				log.Print(err)
				return err
			}
			header.Name = f.Name()
			// See if any header overrides need to be done
			if t.Owner > 0 {
				header.Uid = t.Owner
			}
			if t.Group > 0 {
				header.Gid = t.Group
			}
			if t.FileMode > 0 {
				header.Mode = int64(t.FileMode)
			} else {
				header.Mode = int64(info.Mode().Perm())
			}
			header.ModTime = info.ModTime()
			err = t.Writer.WriteHeader(header)
			if err != nil {
				log.Print(err)
				return err
			}
			_, err = io.Copy(t.Writer, f)
			if err != nil {
				log.Print(err)
				return err
			}
			err = f.Close()
			if err != nil {
				log.Print(err)
				return err
			}
		}
		return nil
	}()
	return &wg, nil
}

// Delete is not implemented
func (t *Tar) Delete() error {
	return nil
}

// Extract extracts the files from src and writes them to the dst.
func (t *Tar) Extract(src io.Reader) error {
	switch t.Format {
	case TarFmt:
		return t.ExtractTar(src)
	case GzipFmt:
		return t.ExtractTgz(src)
	case Bzip2Fmt:
		return t.ExtractTbz(src)
	case LZWFmt:
		return t.ExtractZ(src)
	case LZ4Fmt:
		return t.ExtractLZ4(src)
	default:
		return UnsupportedFmt.NotSupportedError()
	}
	return nil
}

// ExtractTar extracts a tar file using the passed reader
func (t *Tar) ExtractTar(src io.Reader) (err error) {
	tr := tar.NewReader(src)
	for {
		header, err := tr.Next()
		if err != nil {
			if err == io.EOF {
				break
			}
			log.Print(err)
			return err
		}
		fname := header.Name
		// extract is always relative to cwd, for now
		// temporarily commented out because dst is no longer supported
		// TODO add flag for destinatiion
		//fname = filepath.Join(dst, fname)
		fmt.Printf("%s %s\n", fname, strconv.Itoa(int(header.Mode)))
		switch header.Typeflag {
		case tar.TypeDir:
			err = os.MkdirAll(fname, 0744)
			if err != nil {
				log.Print(err)
				return err
			}
			// set the final element to the appropriate permission
			err = os.Chmod(fname, os.FileMode(header.Mode))
			if err != nil {
				log.Print(err)
				return err
			}
		case tar.TypeReg:
			// create the parent directory if necessary
			pdir := filepath.Dir(fname)
			err = os.MkdirAll(pdir, 0744)
			if err != nil {
				log.Print(err)
				return err
			}
			if err != nil {
				log.Print(err)
				return err
			}
			w, err := os.Create(fname)
			if err != nil {
				log.Print(err)
				return err
			}
			io.Copy(w, tr)
			err = os.Chmod(fname, os.FileMode(header.Mode))
			if err != nil {
				log.Print(err)
				return err
			}
			w.Close()
		default:
			err = fmt.Errorf("Unable to extract type: %c in file %s", header.Typeflag, fname)
			log.Print(err)
			return err
		}
	}
	return nil
}

// ExtractGzip reads a GZip using the passed reader.
func (t *Tar) ExtractGzip(src io.Reader) (err error) {
	gR, err := gzip.NewReader(src)
	if err != nil {
		log.Print(err)
		return err
	}
	// Close the file with error handling
	defer func() {
		cerr := gR.Close()
		if cerr != nil && err == nil {
			log.Print(cerr)
			err = cerr
		}
	}()
	err = t.ExtractTar(gR)
	return err
}

// ExtractTgz extracts GZip'd tarballs.
func (t *Tar) ExtractTgz(src io.Reader) error {
	gr, err := gzip.NewReader(src)
	if err != nil {
		log.Print(err)
		return err
	}
	defer gr.Close()
	err = t.ExtractTar(gr)
	return err
}

// ExtractTbz extracts Bzip2 compressed tarballs.
func (t *Tar) ExtractTbz(src io.Reader) error {
	zR := bzip2.NewReader(src)
	return t.ExtractTar(zR)
}

// ExtractZ extracts tarballs compressed with LZW, typically .Z extension.
// TODO fix so that order and width get properly set. Assuming order isn't
// good.
func (t *Tar) ExtractZ(src io.Reader) error {
	zR := lzw.NewReader(src, lzw.LSB, 8)
	defer zR.Close()
	return t.ExtractTar(zR)
}

// ExtractLZ4 extracts LZ4 compressed tarballs.
func (t *Tar) ExtractLZ4(src io.Reader) error {
	lzR := lz4.NewReader(src)
	tr := tar.NewReader(lzR)
	err := t.ExtractTar(tr)
	return err
}

func extractTarFile(hdr *tar.Header, dst string, src io.Reader) error {
	fP := filepath.Join(dst, hdr.Name)
	fI := hdr.FileInfo()
	err := os.MkdirAll(filepath.Dir(fP), fI.Mode())
	if err != nil {
		log.Print(err)
		return err
	}
	if fI.IsDir() {
		return nil
	}
	if fI.Mode()&os.ModeSymlink != 0 {
		return os.Symlink(hdr.Linkname, fP)
	}

	dF, err := os.OpenFile(fP, os.O_RDWR|os.O_CREATE|os.O_TRUNC, fI.Mode())
	if err != nil {
		log.Print(err)
		return err
	}
	defer dF.Close()

	_, err = io.Copy(dF, src)
	if err != nil {
		log.Print(err)
		return err
	}
	return nil
}
<commit_msg>add tar.Extract()<commit_after>package carchivum

// tar implements the tape archive format.

import (
	"archive/tar"
	"compress/bzip2"
	"compress/gzip"
	"compress/lzw"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"time"

	"github.com/MichaelTJones/walk"
	"github.com/pierrec/lz4"
)

// Tar is a struct for a tar, tape archive.
type Tar struct {
	Car
	*tar.Writer
	Format
	sources []string
}

// NewTar returns an initialized Tar struct ready for use.
func NewTar(n string) *Tar {
	return &Tar{Car: Car{Name: n, t0: time.Now()}, Format: defaultFormat, sources: []string{}}
}

// Create creates a tarfile from the passed src('s) and saves it to the dst.
func (t *Tar) Create(src ...string) (cnt int, err error) {
	// If there isn't a destination, return err
	if t.Name == "" {
		err = fmt.Errorf("destination required to create a tar archive")
		log.Print(err)
		return 0, err
	}
	// If there aren't any sources, return err
	if len(src) == 0 {
		err = fmt.Errorf("a source is required to create a tar archive")
		log.Print(err)
		return 0, err
	}
	t.sources = src
	// See if we can create the destination file before processing
	tball, err := os.OpenFile(t.Name, os.O_RDWR|os.O_CREATE, 0744)
	if err != nil {
		log.Print(err)
		return 0, err
	}
	defer func() {
		cerr := tball.Close()
		if cerr != nil && err == nil {
			log.Print(cerr)
			err = cerr
		}
	}()
	switch t.Format {
	case GzipFmt:
		err = t.CreateGzip(tball)
		if err != nil {
			log.Print(err)
			return 0, err
		}
	case Bzip2Fmt:
		err = fmt.Errorf("Bzip2 compression is not supported")
		return 0, err
	case LZWFmt:
		err = t.CreateZ(tball)
		if err != nil {
			log.Print(err)
			return 0, err
		}
	case LZ4Fmt:
		err = t.CreateLZ4(tball)
		if err != nil {
			log.Print(err)
			return 0, err
		}
	default:
		err = fmt.Errorf("Unsupported compression format: %s", t.Format.String())
		return 0, err
	}
	if t.DeleteArchived {
		err := t.removeFiles()
		if err != nil {
			err = fmt.Errorf("an error was encountered while deleting the archived files; some files may not be deleted: %q", err)
			log.Print(err)
			return 0, err
		}
	}
	t.setDelta()
	return int(t.Car.files), nil
}

func (t *Tar) removeFiles() error {
	for _, file := range t.deleteList {
		err := os.Remove(file)
		if err != nil {
			log.Print(err)
			return err
		}
	}
	return nil
}

// CreateGzip creates a GZip using the passed writer.
func (t *Tar) CreateGzip(w io.Writer) (err error) {
	zw := gzip.NewWriter(w)
	// Close the file with error handling
	defer func() {
		cerr := zw.Close()
		if cerr != nil && err == nil {
			log.Print(cerr)
			err = cerr
		}
	}()
	err = t.writeTar(zw)
	return err
}

// CreateLZW compresses using LZW and LSB order using the passed writer.
// TODO: address order so that it doesn't necessarily default to LSB
func (t *Tar) CreateZ(w io.Writer) (err error) {
	zw := lzw.NewWriter(w, lzw.LSB, 8)
	// Close the file with error handling
	defer func() {
		cerr := zw.Close()
		if cerr != nil && err == nil {
			log.Print(cerr)
			err = cerr
		}
	}()
	err = t.writeTar(zw)
	return err
}

// CreateLZ4 creates a LZ4 compressed tarball using the passed writer.
func (t *Tar) CreateLZ4(w io.Writer) (err error) {
	lzW := lz4.NewWriter(w)
	err = t.writeTar(lzW)
	lzW.Close()
	return err
}

func (t *Tar) writeTar(w io.Writer) (err error) {
	t.Writer = tar.NewWriter(w)
	defer func() {
		cerr := t.Writer.Close()
		if cerr != nil && err == nil {
			log.Print(cerr)
			err = cerr
		}
	}()
	t.FileCh = make(chan *os.File)
	wait, err := t.Write()
	if err != nil {
		log.Print(err)
		return err
	}
	var fullPath string
	visitor := func(p string, fi os.FileInfo, err error) error {
		return t.AddFile(fullPath, p, fi, err)
	}
	for _, source := range t.sources {
		fullPath, err = filepath.Abs(source)
		if err != nil {
			log.Print(err)
			return err
		}
		err = walk.Walk(fullPath, visitor)
		if err != nil {
			log.Print(err)
			return err
		}
	}
	close(t.FileCh)
	wait.Wait()
	return err
}

// Write adds the files received on the channel to the tarball.
func (t *Tar) Write() (*sync.WaitGroup, error) {
	var wg sync.WaitGroup
	wg.Add(1)
	go func() error {
		defer wg.Done()
		for f := range t.FileCh {
			info, err := f.Stat()
			if err != nil {
				log.Print(err)
				return err
			}
			if info.IsDir() {
				continue
			}
			header, err := tar.FileInfoHeader(info, "")
			if err != nil {
				log.Print(err)
				return err
			}
			header.Name = f.Name()
			// See if any header overrides need to be done
			if t.Owner > 0 {
				header.Uid = t.Owner
			}
			if t.Group > 0 {
				header.Gid = t.Group
			}
			if t.FileMode > 0 {
				header.Mode = int64(t.FileMode)
			} else {
				header.Mode = int64(info.Mode().Perm())
			}
			header.ModTime = info.ModTime()
			err = t.Writer.WriteHeader(header)
			if err != nil {
				log.Print(err)
				return err
			}
			_, err = io.Copy(t.Writer, f)
			if err != nil {
				log.Print(err)
				return err
			}
			err = f.Close()
			if err != nil {
				log.Print(err)
				return err
			}
		}
		return nil
	}()
	return &wg, nil
}

// Delete is not implemented
func (t *Tar) Delete() error {
	return nil
}

// Extract extracts the files from the src and writes them to the dst. The src
// is either a tar or a compressed tar.
func (t *Tar) Extract() error {
	// find its format
	format, err := getFileFormat(f)
	if err != nil {
		log.Print(err)
		return err
	}
	defer f.Close()
	tar := NewTar(src)
	tar.OutDir = dst
	tar.Format = format
	err = tar.ExtractArchive(f)
	if err != nil {
		log.Print(err)
	}
	return err
}

func (t *Tar) ExtractArchive(src io.Reader) error {
	switch t.Format {
	case TarFmt:
		return t.ExtractTar(src)
	case GzipFmt:
		return t.ExtractTgz(src)
	case Bzip2Fmt:
		return t.ExtractTbz(src)
	case LZWFmt:
		return t.ExtractZ(src)
	case LZ4Fmt:
		return t.ExtractLZ4(src)
	default:
		return UnsupportedFmt.NotSupportedError()
	}
	return nil
}

// ExtractTar extracts a tar file using the passed reader
func (t *Tar) ExtractTar(src io.Reader) (err error) {
	tr := tar.NewReader(src)
	for {
		header, err := tr.Next()
		if err != nil {
			if err == io.EOF {
				break
			}
			log.Print(err)
			return err
		}
		fname := header.Name
		// extract is always relative to cwd, for now
		// temporarily commented out because dst is no longer supported
		// TODO add flag for destinatiion
		//fname = filepath.Join(dst, fname)
		fmt.Printf("%s %s\n", fname, strconv.Itoa(int(header.Mode)))
		switch header.Typeflag {
		case tar.TypeDir:
			err = os.MkdirAll(fname, 0744)
			if err != nil {
				log.Print(err)
				return err
			}
			// set the final element to the appropriate permission
			err = os.Chmod(fname, os.FileMode(header.Mode))
			if err != nil {
				log.Print(err)
				return err
			}
		case tar.TypeReg:
			// create the parent directory if necessary
			pdir := filepath.Dir(fname)
			err = os.MkdirAll(pdir, 0744)
			if err != nil {
				log.Print(err)
				return err
			}
			if err != nil {
				log.Print(err)
				return err
			}
			w, err := os.Create(fname)
			if err != nil {
				log.Print(err)
				return err
			}
			io.Copy(w, tr)
			err = os.Chmod(fname, os.FileMode(header.Mode))
			if err != nil {
				log.Print(err)
				return err
			}
			w.Close()
		default:
			err = fmt.Errorf("Unable to extract type: %c in file %s", header.Typeflag, fname)
			log.Print(err)
			return err
		}
	}
	return nil
}

// ExtractGzip reads a GZip using the passed reader.
func (t *Tar) ExtractGzip(src io.Reader) (err error) {
	gR, err := gzip.NewReader(src)
	if err != nil {
		log.Print(err)
		return err
	}
	// Close the file with error handling
	defer func() {
		cerr := gR.Close()
		if cerr != nil && err == nil {
			log.Print(cerr)
			err = cerr
		}
	}()
	err = t.ExtractTar(gR)
	return err
}

// ExtractTgz extracts GZip'd tarballs.
func (t *Tar) ExtractTgz(src io.Reader) error {
	gr, err := gzip.NewReader(src)
	if err != nil {
		log.Print(err)
		return err
	}
	defer gr.Close()
	err = t.ExtractTar(gr)
	return err
}

// ExtractTbz extracts Bzip2 compressed tarballs.
func (t *Tar) ExtractTbz(src io.Reader) error {
	zR := bzip2.NewReader(src)
	return t.ExtractTar(zR)
}

// ExtractZ extracts tarballs compressed with LZW, typically .Z extension.
// TODO fix so that order and width get properly set. Assuming order isn't
// good.
func (t *Tar) ExtractZ(src io.Reader) error {
	zR := lzw.NewReader(src, lzw.LSB, 8)
	defer zR.Close()
	return t.ExtractTar(zR)
}

// ExtractLZ4 extracts LZ4 compressed tarballs.
func (t *Tar) ExtractLZ4(src io.Reader) error {
	lzR := lz4.NewReader(src)
	tr := tar.NewReader(lzR)
	err := t.ExtractTar(tr)
	return err
}

func extractTarFile(hdr *tar.Header, dst string, src io.Reader) error {
	fP := filepath.Join(dst, hdr.Name)
	fI := hdr.FileInfo()
	err := os.MkdirAll(filepath.Dir(fP), fI.Mode())
	if err != nil {
		log.Print(err)
		return err
	}
	if fI.IsDir() {
		return nil
	}
	if fI.Mode()&os.ModeSymlink != 0 {
		return os.Symlink(hdr.Linkname, fP)
	}

	dF, err := os.OpenFile(fP, os.O_RDWR|os.O_CREATE|os.O_TRUNC, fI.Mode())
	if err != nil {
		log.Print(err)
		return err
	}
	defer dF.Close()

	_, err = io.Copy(dF, src)
	if err != nil {
		log.Print(err)
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package next

import (
	"bufio"
	"encoding/binary"
	"errors"
	"io"
	"log"
	"net"
	"os"
	"reflect"
	"runtime"
)

type Tcp struct {
	Conn   map[string]*net.TCPConn
	Config *Config
	Logger *log.Logger
	routes *Routes
}

const (
	TcpHead       = 0xAA
	TcpTail       = 0x0A
	TcpMaxContent = 2 << 10
)

func NewTcp() *Tcp {
	tcp := &Tcp{
		Conn:   make(map[string]*net.TCPConn),
		Config: NewConfig(),
		Logger: log.New(os.Stdout, "", log.Ldate|log.Ltime),
		routes: NewRoutes(),
	}

	// Load default config if exists
	file := "config.json"
	if fileExists(file) {
		tcp.Config.Load(file)
	}

	return tcp
}

// Pack is a utility function to read from the supplied Writer
// according to the Next protocol spec:
//
//    AA[x][x][x][x][x][x][x][x]55
//      |  (int32) || (binary)
//      |  4-byte  || N-byte
//      ------------------------...
//        size       data
func (t *Tcp) Pack(w io.Writer, data []byte) error {
	// Head
	_, err := w.Write([]byte{TcpHead})
	if err != nil {
		return errors.New("write head fail")
	}

	// Size
	err = binary.Write(w, binary.LittleEndian, int32(len(data))+6)
	if err != nil {
		return errors.New("write size fail")
	}

	// Data
	_, err = w.Write(data)
	if err != nil {
		return errors.New("write data fail")
	}

	// Tail
	_, err = w.Write([]byte{TcpTail})
	if err != nil {
		return errors.New("write tail fail")
	}

	return nil
}

// Unpack is a utility function to read from the supplied Reader
// according to the Next protocol spec:
//
//    AA[x][x][x][x][x][x][x][x]55
//      |  (int32) || (binary)
//      |  4-byte  || N-byte
//      ------------------------...
//        size       data
func (t *Tcp) Unpack(r io.Reader) ([]byte, error) {
	// Check head
	head := make([]byte, 1)
	_, err := r.Read(head)
	if err != nil {
		return nil, err
	}
	if int32(head[0]) != TcpHead {
		return nil, errors.New("data head is error")
	}

	// message size
	var size int32
	err = binary.Read(r, binary.LittleEndian, &size)
	if err != nil {
		return nil, err
	}
	size = size - 6
	if size <= 0 || size > TcpMaxContent {
		return nil, errors.New("data size is error")
	}

	// message binary data
	buf := make([]byte, size)
	_, err = io.ReadFull(r, buf)
	if err != nil {
		return nil, err
	}

	// Check head
	tail := make([]byte, 1)
	_, err = r.Read(tail)
	if err != nil {
		return nil, err
	}
	if int32(tail[0]) != TcpTail {
		return nil, errors.New("data tail is error")
	}

	return buf, nil
}

func (t *Tcp) handler(conn *net.TCPConn, body []byte) {
	ctx := TcpContext{
		Params: make(map[string]string),
		Tcp:    t,
		Fd:     t.Fd(conn),
		conn:   conn,
	}

	// Read json body
	json := NewJson()
	json.Load(body)
	requestPath := json.Get("method").MustString()

	route := t.routes.Match(requestPath, "VIA")
	if route == nil {
		ctx.WriteJSON("404", "request method not found")
		return
	}

	data, error := json.Get("data").Map()
	if error != nil {
		ctx.WriteJSON("403", "data is not map[string]string")
		return
	}
	if len(data) > 0 {
		for key, val := range data {
			v, err := val.(string)
			if !err {
				ctx.WriteJSON("403", "data is not map[string]string")
				return
			}
			ctx.Params[key] = v
		}
	}

	cr := route.cr

	var args []reflect.Value
	handlerType := route.handler.Type()
	if requiresTcpContext(handlerType) {
		args = append(args, reflect.ValueOf(&ctx))
	}
	t.Logger.Print(args)

	match := cr.FindStringSubmatch(requestPath)
	for _, arg := range match[1:] {
		args = append(args, reflect.ValueOf(arg))
	}

	_, err := t.safelyCall(route.handler, args)
	if err != nil {
		//there was an error or panic while calling the handler
		ctx.WriteJSON("500", "Server Error")
	}
}

// Get the integer Unix file descriptor referencing the open file
func (t *Tcp) Fd(conn *net.TCPConn) string {
	return conn.RemoteAddr().String()
}

func (t *Tcp) Pipe(conn *net.TCPConn) {
	fd := t.Fd(conn)
	defer func() {
		t.Logger.Printf("disconnected: %s\n", fd)
		conn.Close()
		delete(t.Conn, fd)
	}()

	// Save in map
	t.Conn[fd] = conn

	// Read data
	reader := bufio.NewReader(conn)
	for {
		body, err := t.Unpack(reader)
		if err != nil {
			if err == io.EOF {
				return
			}
			t.Logger.Print(err)
		}

		t.Logger.Printf("%v", body)
		// Filter heart pack
		if string(body) != "hello" {
			t.handler(conn, body)
		}
		reader.Reset(conn)
	}
}

func (t *Tcp) Run(addr string) {
	tcpAddr, _ := net.ResolveTCPAddr("tcp", addr)
	tcpListener, _ := net.ListenTCP("tcp", tcpAddr)
	defer tcpListener.Close()

	t.Logger.Printf("next tcp serving %s\n", addr)
	for {
		tcpConn, err := tcpListener.AcceptTCP()
		if err != nil {
			continue
		}

		t.Logger.Printf("connected: %s\n", tcpConn.RemoteAddr().String())
		go t.Pipe(tcpConn)
	}
}

// Post adds a handler for the 'Via' TCP method for tcp.
func (t *Tcp) Via(route string, handler interface{}) {
	t.routes.Add(route, "VIA", handler)
}

func (t *Tcp) WriteJSON(conn *net.TCPConn, code, msg string, data ...interface{}) {
	json := NewJson()
	json.Set("code", code)
	json.Set("msg", msg)

	if len(data) > 0 {
		json.Set("data", data[0])
	}
	out, _ := json.Encode()

	t.Pack(conn, out)
}

// safelyCall invokes `function` in recover block
func (t *Tcp) safelyCall(function reflect.Value, args []reflect.Value) (resp []reflect.Value, e interface{}) {
	defer func() {
		if err := recover(); err != nil {
			if t.Config.Bool("panic") {
				// go back to panic
				panic(err)
			} else {
				e = err
				resp = nil
				t.Logger.Println("Handler crashed with error", err)
				for i := 1; ; i += 1 {
					_, file, line, ok := runtime.Caller(i)
					if !ok {
						break
					}
					t.Logger.Println(file, line)
				}
			}
		}
	}()
	return function.Call(args), nil
}

// requiresContext determines whether 'handlerType' contains
// an argument to 'web.Ctx' as its first argument
func requiresTcpContext(handlerType reflect.Type) bool {
	//if the method doesn't take arguments, no
	if handlerType.NumIn() == 0 {
		return false
	}

	//if the first argument is not a pointer, no
	a0 := handlerType.In(0)
	if a0.Kind() != reflect.Ptr {
		return false
	}
	//if the first argument is a context, yes
	if a0.Elem() == reflect.TypeOf(TcpContext{}) {
		return true
	}

	return false
}

// --------
// Tcp Context
// --------
type TcpContext struct {
	Params map[string]string
	Tcp    *Tcp
	Fd     string
	conn   *net.TCPConn
}

// WriteJSON writes json data into the response object.
func (ctx *TcpContext) WriteJSON(code, msg string, data ...interface{}) {
	json := NewJson()
	if ctx.Tcp.Config.Bool("debug.profiler") {
		json.Set("debug", ctx.Params)
	}

	// Back client request version
	if seq, ok := ctx.Params["seq"]; ok {
		json.Set("seq", seq)
	}

	json.Set("code", code)
	json.Set("msg", msg)

	if len(data) > 0 {
		json.Set("data", data[0])
	}
	out, _ := json.Encode()

	ctx.Tcp.Pack(ctx.conn, out)
}
<commit_msg>add reqeust url back to<commit_after>package next

import (
	"bufio"
	"encoding/binary"
	"errors"
	"io"
	"log"
	"net"
	"os"
	"reflect"
	"runtime"
)

type Tcp struct {
	Conn   map[string]*net.TCPConn
	Config *Config
	Logger *log.Logger
	routes *Routes
}

const (
	TcpHead       = 0xAA
	TcpTail       = 0x0A
	TcpMaxContent = 2 << 10
)

func NewTcp() *Tcp {
	tcp := &Tcp{
		Conn:   make(map[string]*net.TCPConn),
		Config: NewConfig(),
		Logger: log.New(os.Stdout, "", log.Ldate|log.Ltime),
		routes: NewRoutes(),
	}

	// Load default config if exists
	file := "config.json"
	if fileExists(file) {
		tcp.Config.Load(file)
	}

	return tcp
}

// Pack is a utility function to read from the supplied Writer
// according to the Next protocol spec:
//
//    AA[x][x][x][x][x][x][x][x]55
//      |  (int32) || (binary)
//      |  4-byte  || N-byte
//      ------------------------...
//        size       data
func (t *Tcp) Pack(w io.Writer, data []byte) error {
	// Head
	_, err := w.Write([]byte{TcpHead})
	if err != nil {
		return errors.New("write head fail")
	}

	// Size
	err = binary.Write(w, binary.LittleEndian, int32(len(data))+6)
	if err != nil {
		return errors.New("write size fail")
	}

	// Data
	_, err = w.Write(data)
	if err != nil {
		return errors.New("write data fail")
	}

	// Tail
	_, err = w.Write([]byte{TcpTail})
	if err != nil {
		return errors.New("write tail fail")
	}

	return nil
}

// Unpack is a utility function to read from the supplied Reader
// according to the Next protocol spec:
//
//    AA[x][x][x][x][x][x][x][x]55
//      |  (int32) || (binary)
//      |  4-byte  || N-byte
//      ------------------------...
//        size       data
func (t *Tcp) Unpack(r io.Reader) ([]byte, error) {
	// Check head
	head := make([]byte, 1)
	_, err := r.Read(head)
	if err != nil {
		return nil, err
	}
	if int32(head[0]) != TcpHead {
		return nil, errors.New("data head is error")
	}

	// message size
	var size int32
	err = binary.Read(r, binary.LittleEndian, &size)
	if err != nil {
		return nil, err
	}
	size = size - 6
	if size <= 0 || size > TcpMaxContent {
		return nil, errors.New("data size is error")
	}

	// message binary data
	buf := make([]byte, size)
	_, err = io.ReadFull(r, buf)
	if err != nil {
		return nil, err
	}

	// Check head
	tail := make([]byte, 1)
	_, err = r.Read(tail)
	if err != nil {
		return nil, err
	}
	if int32(tail[0]) != TcpTail {
		return nil, errors.New("data tail is error")
	}

	return buf, nil
}

func (t *Tcp) handler(conn *net.TCPConn, body []byte) {
	ctx := TcpContext{
		Params: make(map[string]string),
		Tcp:    t,
		Fd:     t.Fd(conn),
		conn:   conn,
	}

	// Read json body
	json := NewJson()
	json.Load(body)
	requestPath := json.Get("method").MustString()

	route := t.routes.Match(requestPath, "VIA")
	if route == nil {
		ctx.WriteJSON("404", "request method not found")
		return
	}
	ctx.Params["method"] = requestPath

	data, error := json.Get("data").Map()
	if error != nil {
		ctx.WriteJSON("403", "data is not map[string]string")
		return
	}
	if len(data) > 0 {
		for key, val := range data {
			v, err := val.(string)
			if !err {
				ctx.WriteJSON("403", "data is not map[string]string")
				return
			}
			ctx.Params[key] = v
		}
	}

	cr := route.cr

	var args []reflect.Value
	handlerType := route.handler.Type()
	if requiresTcpContext(handlerType) {
		args = append(args, reflect.ValueOf(&ctx))
	}
	t.Logger.Print(args)

	match := cr.FindStringSubmatch(requestPath)
	for _, arg := range match[1:] {
		args = append(args, reflect.ValueOf(arg))
	}

	_, err := t.safelyCall(route.handler, args)
	if err != nil {
		//there was an error or panic while calling the handler
		ctx.WriteJSON("500", "Server Error")
	}
}

// Get the integer Unix file descriptor referencing the open file
func (t *Tcp) Fd(conn *net.TCPConn) string {
	return conn.RemoteAddr().String()
}

func (t *Tcp) Pipe(conn *net.TCPConn) {
	fd := t.Fd(conn)
	defer func() {
		t.Logger.Printf("disconnected: %s\n", fd)
		conn.Close()
		delete(t.Conn, fd)
	}()

	// Save in map
	t.Conn[fd] = conn

	// Read data
	reader := bufio.NewReader(conn)
	for {
		body, err := t.Unpack(reader)
		if err != nil {
			if err == io.EOF {
				return
			}
			t.Logger.Print(err)
		}

		t.Logger.Printf("%v", body)
		// Filter heart pack
		if string(body) != "hello" {
			t.handler(conn, body)
		}
		reader.Reset(conn)
	}
}

func (t *Tcp) Run(addr string) {
	tcpAddr, _ := net.ResolveTCPAddr("tcp", addr)
	tcpListener, _ := net.ListenTCP("tcp", tcpAddr)
	defer tcpListener.Close()

	t.Logger.Printf("next tcp serving %s\n", addr)
	for {
		tcpConn, err := tcpListener.AcceptTCP()
		if err != nil {
			continue
		}

		t.Logger.Printf("connected: %s\n", tcpConn.RemoteAddr().String())
		go t.Pipe(tcpConn)
	}
}

// Post adds a handler for the 'Via' TCP method for tcp.
func (t *Tcp) Via(route string, handler interface{}) {
	t.routes.Add(route, "VIA", handler)
}

func (t *Tcp) WriteJSON(conn *net.TCPConn, code, msg string, data ...interface{}) {
	json := NewJson()
	json.Set("code", code)
	json.Set("msg", msg)

	if len(data) > 0 {
		json.Set("data", data[0])
	}
	out, _ := json.Encode()

	t.Pack(conn, out)
}

// safelyCall invokes `function` in recover block
func (t *Tcp) safelyCall(function reflect.Value, args []reflect.Value) (resp []reflect.Value, e interface{}) {
	defer func() {
		if err := recover(); err != nil {
			if t.Config.Bool("panic") {
				// go back to panic
				panic(err)
			} else {
				e = err
				resp = nil
				t.Logger.Println("Handler crashed with error", err)
				for i := 1; ; i += 1 {
					_, file, line, ok := runtime.Caller(i)
					if !ok {
						break
					}
					t.Logger.Println(file, line)
				}
			}
		}
	}()
	return function.Call(args), nil
}

// requiresContext determines whether 'handlerType' contains
// an argument to 'web.Ctx' as its first argument
func requiresTcpContext(handlerType reflect.Type) bool {
	//if the method doesn't take arguments, no
	if handlerType.NumIn() == 0 {
		return false
	}

	//if the first argument is not a pointer, no
	a0 := handlerType.In(0)
	if a0.Kind() != reflect.Ptr {
		return false
	}
	//if the first argument is a context, yes
	if a0.Elem() == reflect.TypeOf(TcpContext{}) {
		return true
	}

	return false
}

// --------
// Tcp Context
// --------
type TcpContext struct {
	Params map[string]string
	Tcp    *Tcp
	Fd     string
	conn   *net.TCPConn
}

// WriteJSON writes json data into the response object.
func (ctx *TcpContext) WriteJSON(code, msg string, data ...interface{}) {
	json := NewJson()
	if ctx.Tcp.Config.Bool("debug.profiler") {
		json.Set("debug", ctx.Params)
	}

	// Back client request version
	if seq, ok := ctx.Params["seq"]; ok {
		json.Set("seq", seq)
	}
	// Back request method
	if method, ok := ctx.Params["method"]; ok {
		json.Set("method", method)
	}

	json.Set("code", code)
	json.Set("msg", msg)

	if len(data) > 0 {
		json.Set("data", data[0])
	}
	out, _ := json.Encode()

	ctx.Tcp.Pack(ctx.conn, out)
}
<|endoftext|>


<commit_before>// Copyright 2016 The Prometheus 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.

package main

import (
	"bufio"
	"crypto/tls"
	"fmt"
	"net"
	"regexp"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/config"
	"github.com/prometheus/common/log"
)

func dialTCP(target string, module Module, registry *prometheus.Registry) (net.Conn, error) {
	var dialProtocol, dialTarget string
	dialer := &net.Dialer{Timeout: module.Timeout}

	targetAddress, port, err := net.SplitHostPort(target)
	if err != nil {
		return nil, err
	}

	ip, err := chooseProtocol(module.TCP.PreferredIPProtocol, targetAddress, registry)
	if err != nil {
		return nil, err
	}

	if ip.IP.To4() == nil {
		dialProtocol = "tcp6"
	} else {
		dialProtocol = "tcp4"
	}
	dialTarget = net.JoinHostPort(ip.String(), port)

	if !module.TCP.TLS {
		return dialer.Dial(dialProtocol, dialTarget)
	}
	tlsConfig, err := config.NewTLSConfig(&module.TCP.TLSConfig)
	if err != nil {
		return nil, err
	}
	return tls.DialWithDialer(dialer, dialProtocol, dialTarget, tlsConfig)
}

func probeTCP(target string, module Module, registry *prometheus.Registry) bool {
	probeSSLEarliestCertExpiry := prometheus.NewGauge(prometheus.GaugeOpts{
		Name: "probe_ssl_earliest_cert_expiry",
		Help: "Returns earliest SSL cert expiry date",
	})
	registry.MustRegister(probeSSLEarliestCertExpiry)
	deadline := time.Now().Add(module.Timeout)
	conn, err := dialTCP(target, module, registry)
	if err != nil {
		return false
	}
	defer conn.Close()

	// Set a deadline to prevent the following code from blocking forever.
	// If a deadline cannot be set, better fail the probe by returning an error
	// now rather than blocking forever.
	if err := conn.SetDeadline(deadline); err != nil {
		return false
	}
	if module.TCP.TLS {
		state := conn.(*tls.Conn).ConnectionState()
		probeSSLEarliestCertExpiry.Set(float64(getEarliestCertExpiry(&state).UnixNano()) / 1e9)
	}
	scanner := bufio.NewScanner(conn)
	for _, qr := range module.TCP.QueryResponse {
		log.Debugf("Processing query response entry %+v", qr)
		send := qr.Send
		if qr.Expect != "" {
			re, err := regexp.Compile(qr.Expect)
			if err != nil {
				log.Errorf("Could not compile %q into regular expression: %v", qr.Expect, err)
				return false
			}
			var match []int
			// Read lines until one of them matches the configured regexp.
			for scanner.Scan() {
				log.Debugf("read %q\n", scanner.Text())
				match = re.FindSubmatchIndex(scanner.Bytes())
				if match != nil {
					log.Debugf("regexp %q matched %q", re, scanner.Text())
					break
				}
			}
			if scanner.Err() != nil {
				return false
			}
			if match == nil {
				return false
			}
			send = string(re.Expand(nil, []byte(send), scanner.Bytes(), match))
		}
		if send != "" {
			log.Debugf("Sending %q", send)
			if _, err := fmt.Fprintf(conn, "%s\n", send); err != nil {
				return false
			}
		}
	}
	return true
}
<commit_msg>Log tcp errors (#188)<commit_after>// Copyright 2016 The Prometheus 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.

package main

import (
	"bufio"
	"crypto/tls"
	"fmt"
	"net"
	"regexp"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/config"
	"github.com/prometheus/common/log"
)

func dialTCP(target string, module Module, registry *prometheus.Registry) (net.Conn, error) {
	var dialProtocol, dialTarget string
	dialer := &net.Dialer{Timeout: module.Timeout}

	targetAddress, port, err := net.SplitHostPort(target)
	if err != nil {
		log.Errorf("Error splitting target address and port: %v", err)
		return nil, err
	}

	ip, err := chooseProtocol(module.TCP.PreferredIPProtocol, targetAddress, registry)
	if err != nil {
		log.Errorf("Error choosing protocol: %v", err)
		return nil, err
	}

	if ip.IP.To4() == nil {
		dialProtocol = "tcp6"
	} else {
		dialProtocol = "tcp4"
	}
	dialTarget = net.JoinHostPort(ip.String(), port)

	if !module.TCP.TLS {
		return dialer.Dial(dialProtocol, dialTarget)
	}
	tlsConfig, err := config.NewTLSConfig(&module.TCP.TLSConfig)
	if err != nil {
		log.Errorf("Error creating TLS configuration: %v", err)
		return nil, err
	}
	return tls.DialWithDialer(dialer, dialProtocol, dialTarget, tlsConfig)
}

func probeTCP(target string, module Module, registry *prometheus.Registry) bool {
	probeSSLEarliestCertExpiry := prometheus.NewGauge(prometheus.GaugeOpts{
		Name: "probe_ssl_earliest_cert_expiry",
		Help: "Returns earliest SSL cert expiry date",
	})
	registry.MustRegister(probeSSLEarliestCertExpiry)
	deadline := time.Now().Add(module.Timeout)
	conn, err := dialTCP(target, module, registry)
	if err != nil {
		log.Errorf("Error dialing TCP: %v", err)
		return false
	}
	defer conn.Close()

	// Set a deadline to prevent the following code from blocking forever.
	// If a deadline cannot be set, better fail the probe by returning an error
	// now rather than blocking forever.
	if err := conn.SetDeadline(deadline); err != nil {
		log.Errorf("Error setting deadline: %v", err)
		return false
	}
	if module.TCP.TLS {
		state := conn.(*tls.Conn).ConnectionState()
		probeSSLEarliestCertExpiry.Set(float64(getEarliestCertExpiry(&state).UnixNano()) / 1e9)
	}
	scanner := bufio.NewScanner(conn)
	for _, qr := range module.TCP.QueryResponse {
		log.Debugf("Processing query response entry %+v", qr)
		send := qr.Send
		if qr.Expect != "" {
			re, err := regexp.Compile(qr.Expect)
			if err != nil {
				log.Errorf("Could not compile %q into regular expression: %v", qr.Expect, err)
				return false
			}
			var match []int
			// Read lines until one of them matches the configured regexp.
			for scanner.Scan() {
				log.Debugf("read %q\n", scanner.Text())
				match = re.FindSubmatchIndex(scanner.Bytes())
				if match != nil {
					log.Debugf("regexp %q matched %q", re, scanner.Text())
					break
				}
			}
			if scanner.Err() != nil {
				return false
			}
			if match == nil {
				return false
			}
			send = string(re.Expand(nil, []byte(send), scanner.Bytes(), match))
		}
		if send != "" {
			log.Debugf("Sending %q", send)
			if _, err := fmt.Fprintf(conn, "%s\n", send); err != nil {
				return false
			}
		}
	}
	return true
}
<|endoftext|>


<commit_before>package txtdirect

import (
	"bytes"
	"compress/gzip"
	"context"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"sync"
	"time"

	"github.com/cretz/bine/tor"
	"github.com/mholt/caddy"
	cproxy "github.com/mholt/caddy/caddyhttp/proxy"
	"golang.org/x/net/proxy"
	lumberjack "gopkg.in/natefinch/lumberjack.v2"
)

// DefaultOnionServicePort is the port used to serve the onion service on
const DefaultOnionServicePort = 4242

// Tor type config struct
type Tor struct {
	Enable bool
	// Socks5 proxy port
	Port      int
	DataDir   string
	Torrc     string
	DebugMode bool
	LogFile   string

	instance        *tor.Tor
	contextCanceler context.CancelFunc
	onion           *tor.OnionService
}

type TorResponse struct {
	headers    http.Header
	body       []byte
	bodyReader bytes.Buffer
	bodyWriter bytes.Buffer
	status     int
}

// TODO: Discuss these values
const (
	torProxyKeepalive = 30000000
	torFallbackDelay  = 30000000 * time.Millisecond
	torProxyTimeout   = 30000000 * time.Second
)

var bufferPool = sync.Pool{New: createBuffer}

func createBuffer() interface{} {
	return make([]byte, 0, 32*1024)
}

func (t *Tor) Start(c *caddy.Controller) {
	var debugger io.Writer
	if t.DebugMode {
		if t.LogFile != "" {
			debugger = &lumberjack.Logger{
				Filename:   t.LogFile,
				MaxSize:    100,
				MaxAge:     14,
				MaxBackups: 10,
			}
		}
		debugger = os.Stdout
	}

	torInstance, err := tor.Start(nil, &tor.StartConf{
		NoAutoSocksPort: true,
		ExtraArgs:       []string{"--SocksPort", strconv.Itoa(t.Port)},
		TempDataDirBase: t.DataDir,
		TorrcFile:       t.Torrc,
		DebugWriter:     debugger,
	})
	if err != nil {
		log.Panicf("Unable to start Tor: %v", err)
	}

	listenCtx := context.Background()

	onion, err := torInstance.Listen(listenCtx, &tor.ListenConf{LocalPort: 8868, RemotePorts: []int{80}})
	if err != nil {
		log.Panicf("Unable to start onion service: %v", err)
	}

	t.onion = onion
	t.instance = torInstance
}

// Stop stops the tor instance, context listener and the onion service
func (t *Tor) Stop() error {
	if err := t.instance.Close(); err != nil {
		return fmt.Errorf("[txtdirect]: Couldn't close the tor instance. %s", err.Error())
	}
	t.onion.Close()
	return nil
}

// Proxy redirects the request to the local onion serivce and the actual proxying
// happens inside onion service's http handler
func (t *Tor) Proxy(w http.ResponseWriter, r *http.Request, rec record, c Config) error {
	u, err := url.Parse(rec.To)
	if err != nil {
		return err
	}

	// Create a socks5 dialer
	dialer, err := proxy.SOCKS5("tcp", fmt.Sprintf("127.0.0.1:%d", t.Port), nil, proxy.Direct)
	if err != nil {
		log.Fatal(err)
	}

	reverseProxy := cproxy.NewSingleHostReverseProxy(u, "", torProxyKeepalive, torProxyTimeout, torFallbackDelay)
	reverseProxy.Transport = &http.Transport{
		Dial: dialer.Dial,
	}

	tmpResposne := TorResponse{headers: make(http.Header)}
	if err := reverseProxy.ServeHTTP(&tmpResposne, r, nil); err != nil {
		return fmt.Errorf("[txtdirect]: Coudln't proxy the request to the background onion service. %s", err.Error())
	}

	// Decompress the body based on "Content-Encoding" header and write to a writer buffer
	if err := tmpResposne.WriteBody(); err != nil {
		return fmt.Errorf("[txtdirect]: Couldn't write the response body: %s", err.Error())
	}

	// Replace the URL hosts with the request's host
	if err := tmpResposne.ReplaceBody(u.Scheme, u.Host, r.Host); err != nil {
		return fmt.Errorf("[txtdirect]: Couldn't replace urls inside the response body: %s", err.Error())
	}

	copyHeader(w.Header(), tmpResposne.Header())

	// Write the final response from the temporary ResponseWriter to the main ResponseWriter
	if _, err := w.Write(tmpResposne.Body()); err != nil {
		return fmt.Errorf("[txtdirect]: Couldn't write the temporary response to main response body: %s", err.Error())
	}

	return nil
}

// ParseTor parses the txtdirect config for Tor proxy
func (t *Tor) ParseTor(c *caddy.Controller) error {
	switch c.Val() {
	case "port":
		value, err := strconv.Atoi(c.RemainingArgs()[0])
		if err != nil {
			return fmt.Errorf("The given value for port field is not standard. It should be an integer")
		}
		t.Port = value

	case "datadir":
		t.DataDir = c.RemainingArgs()[0]

	case "torrc":
		t.Torrc = c.RemainingArgs()[0]

	case "debug_mode":
		value, err := strconv.ParseBool(c.RemainingArgs()[0])
		if err != nil {
			return fmt.Errorf("The given value for debug_mode field is not standard. It should be a boolean")
		}
		t.DebugMode = value

	case "logfile":
		t.LogFile = c.RemainingArgs()[0]

	default:
		return c.ArgErr() // unhandled option for tor
	}
	return nil
}

// SetDefaults sets the default values for prometheus config
// if the fields are empty
func (t *Tor) SetDefaults() {
	if t.Port == 0 {
		t.Port = DefaultOnionServicePort
	}
}

// Header returns response headers
func (t *TorResponse) Header() http.Header {
	return t.headers
}

func (t *TorResponse) Write(body []byte) (int, error) {
	reader := bytes.NewReader(body)
	pooledIoCopy(&t.bodyReader, reader)
	t.body = body
	return len(body), nil
}

// Body returns response's body. This method should only get called after WriteBody()
func (t *TorResponse) Body() []byte {
	return t.bodyWriter.Bytes()
}

// WriteHeader Writes the given status code to response
func (t *TorResponse) WriteHeader(status int) {
	t.status = status
}

func (t *TorResponse) ReplaceBody(scheme, to, host string) error {
	replacedBody := bytes.Replace(t.bodyWriter.Bytes(), []byte(scheme+"://"+to), []byte(scheme+"://"+host), -1)
	t.bodyWriter.Reset()
	if _, err := t.bodyWriter.Write(replacedBody); err != nil {
		return err
	}
	return nil
}

func (t *TorResponse) WriteBody() error {
	switch t.Header().Get("Content-Encoding") {
	case "gzip":
		reader, err := gzip.NewReader(&t.bodyReader)
		if err != nil {
			return err
		}
		defer reader.Close()
		_, err = io.Copy(&t.bodyWriter, reader)
		if err != nil {
			return err
		}
		t.Header().Del("Content-Encoding")
	default:
		_, err := io.Copy(&t.bodyWriter, &t.bodyReader)
		if err != nil {
			return err
		}
	}
	return nil
}

var skipHeaders = map[string]struct{}{
	"Content-Type":        {},
	"Content-Disposition": {},
	"Accept-Ranges":       {},
	"Set-Cookie":          {},
	"Cache-Control":       {},
	"Expires":             {},
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		if _, ok := dst[k]; ok {
			if _, shouldSkip := skipHeaders[k]; shouldSkip {
				continue
			}
			if k != "Server" {
				dst.Del(k)
			}
		}
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

func pooledIoCopy(dst io.Writer, src io.Reader) {
	buf := bufferPool.Get().([]byte)
	defer bufferPool.Put(buf)

	bufCap := cap(buf)
	io.CopyBuffer(dst, src, buf[0:bufCap:bufCap])
}
<commit_msg>Overwrite Host<commit_after>package txtdirect

import (
	"bytes"
	"compress/gzip"
	"context"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"sync"
	"time"

	"github.com/cretz/bine/tor"
	"github.com/mholt/caddy"
	cproxy "github.com/mholt/caddy/caddyhttp/proxy"
	"golang.org/x/net/proxy"
	lumberjack "gopkg.in/natefinch/lumberjack.v2"
)

// DefaultOnionServicePort is the port used to serve the onion service on
const DefaultOnionServicePort = 4242

// Tor type config struct
type Tor struct {
	Enable bool
	// Socks5 proxy port
	Port      int
	DataDir   string
	Torrc     string
	DebugMode bool
	LogFile   string

	instance        *tor.Tor
	contextCanceler context.CancelFunc
	onion           *tor.OnionService
}

type TorResponse struct {
	headers    http.Header
	body       []byte
	bodyReader bytes.Buffer
	bodyWriter bytes.Buffer
	status     int
}

// TODO: Discuss these values
const (
	torProxyKeepalive = 30000000
	torFallbackDelay  = 30000000 * time.Millisecond
	torProxyTimeout   = 30000000 * time.Second
)

var bufferPool = sync.Pool{New: createBuffer}

func createBuffer() interface{} {
	return make([]byte, 0, 32*1024)
}

func (t *Tor) Start(c *caddy.Controller) {
	var debugger io.Writer
	if t.DebugMode {
		if t.LogFile != "" {
			debugger = &lumberjack.Logger{
				Filename:   t.LogFile,
				MaxSize:    100,
				MaxAge:     14,
				MaxBackups: 10,
			}
		}
		debugger = os.Stdout
	}

	torInstance, err := tor.Start(nil, &tor.StartConf{
		NoAutoSocksPort: true,
		ExtraArgs:       []string{"--SocksPort", strconv.Itoa(t.Port)},
		TempDataDirBase: t.DataDir,
		TorrcFile:       t.Torrc,
		DebugWriter:     debugger,
	})
	if err != nil {
		log.Panicf("Unable to start Tor: %v", err)
	}

	listenCtx := context.Background()

	onion, err := torInstance.Listen(listenCtx, &tor.ListenConf{LocalPort: 8868, RemotePorts: []int{80}})
	if err != nil {
		log.Panicf("Unable to start onion service: %v", err)
	}

	t.onion = onion
	t.instance = torInstance
}

// Stop stops the tor instance, context listener and the onion service
func (t *Tor) Stop() error {
	if err := t.instance.Close(); err != nil {
		return fmt.Errorf("[txtdirect]: Couldn't close the tor instance. %s", err.Error())
	}
	t.onion.Close()
	return nil
}

// Proxy redirects the request to the local onion serivce and the actual proxying
// happens inside onion service's http handler
func (t *Tor) Proxy(w http.ResponseWriter, r *http.Request, rec record, c Config) error {
	u, err := url.Parse(rec.To)
	if err != nil {
		return err
	}

	// Use Proxied Host
	r.Host = u.Host

	// Create a socks5 dialer
	dialer, err := proxy.SOCKS5("tcp", fmt.Sprintf("127.0.0.1:%d", t.Port), nil, proxy.Direct)
	if err != nil {
		log.Fatal(err)
	}

	reverseProxy := cproxy.NewSingleHostReverseProxy(u, "", torProxyKeepalive, torProxyTimeout, torFallbackDelay)
	reverseProxy.Transport = &http.Transport{
		Dial: dialer.Dial,
	}

	tmpResposne := TorResponse{headers: make(http.Header)}
	if err := reverseProxy.ServeHTTP(&tmpResposne, r, nil); err != nil {
		return fmt.Errorf("[txtdirect]: Coudln't proxy the request to the background onion service. %s", err.Error())
	}

	// Decompress the body based on "Content-Encoding" header and write to a writer buffer
	if err := tmpResposne.WriteBody(); err != nil {
		return fmt.Errorf("[txtdirect]: Couldn't write the response body: %s", err.Error())
	}

	// Replace the URL hosts with the request's host
	if err := tmpResposne.ReplaceBody(u.Scheme, u.Host, r.Host); err != nil {
		return fmt.Errorf("[txtdirect]: Couldn't replace urls inside the response body: %s", err.Error())
	}

	copyHeader(w.Header(), tmpResposne.Header())

	// Write the final response from the temporary ResponseWriter to the main ResponseWriter
	if _, err := w.Write(tmpResposne.Body()); err != nil {
		return fmt.Errorf("[txtdirect]: Couldn't write the temporary response to main response body: %s", err.Error())
	}

	return nil
}

// ParseTor parses the txtdirect config for Tor proxy
func (t *Tor) ParseTor(c *caddy.Controller) error {
	switch c.Val() {
	case "port":
		value, err := strconv.Atoi(c.RemainingArgs()[0])
		if err != nil {
			return fmt.Errorf("The given value for port field is not standard. It should be an integer")
		}
		t.Port = value

	case "datadir":
		t.DataDir = c.RemainingArgs()[0]

	case "torrc":
		t.Torrc = c.RemainingArgs()[0]

	case "debug_mode":
		value, err := strconv.ParseBool(c.RemainingArgs()[0])
		if err != nil {
			return fmt.Errorf("The given value for debug_mode field is not standard. It should be a boolean")
		}
		t.DebugMode = value

	case "logfile":
		t.LogFile = c.RemainingArgs()[0]

	default:
		return c.ArgErr() // unhandled option for tor
	}
	return nil
}

// SetDefaults sets the default values for prometheus config
// if the fields are empty
func (t *Tor) SetDefaults() {
	if t.Port == 0 {
		t.Port = DefaultOnionServicePort
	}
}

// Header returns response headers
func (t *TorResponse) Header() http.Header {
	return t.headers
}

func (t *TorResponse) Write(body []byte) (int, error) {
	reader := bytes.NewReader(body)
	pooledIoCopy(&t.bodyReader, reader)
	t.body = body
	return len(body), nil
}

// Body returns response's body. This method should only get called after WriteBody()
func (t *TorResponse) Body() []byte {
	return t.bodyWriter.Bytes()
}

// WriteHeader Writes the given status code to response
func (t *TorResponse) WriteHeader(status int) {
	t.status = status
}

func (t *TorResponse) ReplaceBody(scheme, to, host string) error {
	replacedBody := bytes.Replace(t.bodyWriter.Bytes(), []byte(scheme+"://"+to), []byte(scheme+"://"+host), -1)
	t.bodyWriter.Reset()
	if _, err := t.bodyWriter.Write(replacedBody); err != nil {
		return err
	}
	return nil
}

func (t *TorResponse) WriteBody() error {
	switch t.Header().Get("Content-Encoding") {
	case "gzip":
		reader, err := gzip.NewReader(&t.bodyReader)
		if err != nil {
			return err
		}
		defer reader.Close()
		_, err = io.Copy(&t.bodyWriter, reader)
		if err != nil {
			return err
		}
		t.Header().Del("Content-Encoding")
	default:
		_, err := io.Copy(&t.bodyWriter, &t.bodyReader)
		if err != nil {
			return err
		}
	}
	return nil
}

var skipHeaders = map[string]struct{}{
	"Content-Type":        {},
	"Content-Disposition": {},
	"Accept-Ranges":       {},
	"Set-Cookie":          {},
	"Cache-Control":       {},
	"Expires":             {},
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		if _, ok := dst[k]; ok {
			if _, shouldSkip := skipHeaders[k]; shouldSkip {
				continue
			}
			if k != "Server" {
				dst.Del(k)
			}
		}
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

func pooledIoCopy(dst io.Writer, src io.Reader) {
	buf := bufferPool.Get().([]byte)
	defer bufferPool.Put(buf)

	bufCap := cap(buf)
	io.CopyBuffer(dst, src, buf[0:bufCap:bufCap])
}
<|endoftext|>