<commit_before>package whproxy

import (
	"bytes"
	"io"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"
	"time"

	"sync"

	"github.com/gorilla/websocket"
	log "github.com/sirupsen/logrus"
	"github.com/taskcluster/webhooktunnel/util"
	"github.com/taskcluster/webhooktunnel/wsmux"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  64 * 1024,
	WriteBufferSize: 64 * 1024,
}

func genLogger(fname string) *log.Logger {
	file, err := os.Create(fname)
	if err != nil {
		panic(err)
	}
	logger := &log.Logger{
		Out:       file,
		Formatter: new(log.TextFormatter),
		Level:     log.DebugLevel,
	}
	return logger
}

func TestProxyRegister(t *testing.T) {
	//  start proxy server
	proxy := New(Config{Upgrader: upgrader, Logger: genLogger("register-test")})
	server := httptest.NewServer(proxy.GetHandler())
	defer server.Close()

	// get url
	wsURL := util.MakeWsURL(server.URL)

	// create address to dial
	workerID := "validWorkerID"
	dialAddr := wsURL + "/register/" + workerID

	// dial connection to proxy
	conn1, _, err := websocket.DefaultDialer.Dial(dialAddr, nil)
	if err != nil {
		t.Fatal(err)
	}

	defer func() {
		_ = conn1.Close()
	}()
	// second connection should fail
	conn2, _, err := websocket.DefaultDialer.Dial(dialAddr, nil)
	if err == nil {
		defer func() {
			_ = conn2.Close()
		}()
		t.Fatalf("bad status code: connection should fail")
	}
}

// TestProxyRequest
func TestProxyRequest(t *testing.T) {
	proxy := New(Config{Upgrader: upgrader, Logger: genLogger("request-test")})
	server := httptest.NewServer(proxy.GetHandler())
	defer server.Close()

	// get url
	wsURL := util.MakeWsURL(server.URL)
	// makeshift client
	clientWs, _, err := websocket.DefaultDialer.Dial(wsURL+"/register/workerID", nil)
	if err != nil {
		t.Fatal(err)
	}

	// handler to serve client requests
	clientHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.Method {
		case http.MethodGet:
			_, _ = w.Write([]byte("GET successful"))
		case http.MethodPost:
			_, _ = io.Copy(w, r.Body)
		default:
			http.NotFound(w, r)
		}
	})

	// serve client endpoint
	clientServer := &http.Server{Handler: clientHandler}
	go func() {
		_ = clientServer.Serve(wsmux.Client(clientWs, wsmux.Config{}))
	}()
	defer func() {
		_ = clientServer.Close()
	}()

	// make requests
	viewer := &http.Client{}
	servURL := server.URL

	// GET request
	resp, err := viewer.Get(servURL + "/workerID/")
	if err != nil {
		t.Fatal(err)
	}
	if resp.StatusCode != 200 {
		t.Log(resp)
		t.Fatalf("bad status code on get request")
	}
	reply, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(reply, []byte("GET successful")) {
		t.Fatalf("GET failed. Bad message")
	}

	// POST request
	resp, err = viewer.Post(servURL+"/workerID/", "application/text", bytes.NewBuffer([]byte("message")))
	if err != nil {
		t.Fatal(err)
	}
	if resp.StatusCode != 200 {
		t.Fatalf("bad status code on post request")
	}
	reply, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(reply, []byte("message")) {
		t.Fatalf("POST failed. Bad message")
	}

	// GET request to invalid id
	resp, err = viewer.Get(servURL + "/notWorkerID/")
	if resp.StatusCode != 404 {
		t.Fatalf("request should fail with 404")
	}
}

func TestProxyWebsocket(t *testing.T) {
	proxy := New(Config{Upgrader: upgrader})
	server := httptest.NewServer(proxy.GetHandler())
	wsURL := util.MakeWsURL(server.URL)
	defer server.Close()

	clientHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !websocket.IsWebSocketUpgrade(r) {
			http.NotFound(w, r)
		}

		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			t.Fatal(err)
		}

		mt, buf, err := conn.ReadMessage()
		if err != nil {
			t.Fatal(err)
		}

		err = conn.WriteMessage(mt, buf)
		if err != nil {
			t.Fatal(err)
		}
	})

	// register worker and serve http
	clientWs, _, err := websocket.DefaultDialer.Dial(wsURL+"/register/wsWorker/", nil)
	if err != nil {
		t.Fatal(err)
	}

	clientServer := &http.Server{Handler: clientHandler}
	go func() {
		_ = clientServer.Serve(wsmux.Client(clientWs, wsmux.Config{}))
	}()
	defer func() {
		_ = clientServer.Close()
	}()

	// create websocket connection
	conn, _, err := websocket.DefaultDialer.Dial(wsURL+"/wsWorker/", nil)
	if err != nil {
		t.Fatal(err)
	}
	defer func() {
		_ = conn.Close()
	}()

	// Generate 1M message
	message := make([]byte, 0)
	for i := 0; i < 1024*1024; i++ {
		message = append(message, byte(i%127))
	}

	err = conn.WriteMessage(websocket.BinaryMessage, message)
	if err != nil {
		t.Fatal(err)
	}

	_, buf, err := conn.ReadMessage()
	if !bytes.Equal(buf, message) {
		t.Fatalf("websocket test failed. Bad message")
	}
}

// ensure control messages are proxied
func TestWebsocketProxyControl(t *testing.T) {
	logger := genLogger("ws-control-test")
	proxy := New(Config{Upgrader: upgrader})
	//serve proxy
	server := httptest.NewServer(proxy.GetHandler())
	wsURL := util.MakeWsURL(server.URL)
	defer server.Close()

	// mechanism to know test has completed
	var wg sync.WaitGroup
	wg.Add(2)
	done := func() chan bool {
		tdone := make(chan bool, 1)
		go func() {
			wg.Wait()
			close(tdone)
		}()
		return tdone
	}

	clientHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !websocket.IsWebSocketUpgrade(r) {
			http.NotFound(w, r)
		}

		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			t.Fatal(err)
		}

		// set ping handler. Decrement wg to ensure ping frame was received
		conn.SetPingHandler(func(appData string) error {
			defer wg.Done()
			return conn.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(500*time.Millisecond))
		})

		// Read message to make sure ping was received
		for {
			_, _, err = conn.NextReader()
			if err != nil {
				break
			}
		}
	})

	// register worker and serve http
	clientWs, _, err := websocket.DefaultDialer.Dial(wsURL+"/register/wsWorker/", nil)
	if err != nil {
		t.Fatal(err)
	}

	clientServer := &http.Server{Handler: clientHandler}
	go func() {
		_ = clientServer.Serve(wsmux.Client(clientWs, wsmux.Config{}))
	}()
	defer func() {
		_ = clientServer.Close()
	}()

	// create websocket connection
	conn, _, err := websocket.DefaultDialer.Dial(wsURL+"/wsWorker/", nil)
	if err != nil {
		t.Fatal(err)
	}
	defer func() {
		_ = conn.Close()
	}()

	// ****************************
	// Set Pong Handler. Decrement wg when pong handler fires to ensure that
	// pong is called
	conn.SetPongHandler(func(appData string) error {
		defer wg.Done()
		logger.Printf("received pong: %s", appData)
		if appData != "ping" {
			t.Fatal("bad pong")
		}
		return nil
	})

	// set timer for timing out test
	timer := time.NewTimer(3 * time.Second)

	// start reading messages to ensure pong is received
	go func() {
		for {
			_, _, err = conn.NextReader()
			if err != nil {
				break
			}
		}
	}()

	err = conn.WriteControl(websocket.PingMessage, []byte("ping"), time.Now().Add(1*time.Second))
	if err != nil {
		t.Fatal(err)
	}

	select {
	case <-timer.C:
		t.Fatalf("test failed: timeout")
	case <-done():
	}

}

// Ensure websocket close is proxied
func TestWebSocketClosure(t *testing.T) {
	logger := genLogger("ws-closure-test")
	proxy := New(Config{Upgrader: upgrader})
	//serve proxy
	server := httptest.NewServer(proxy.GetHandler())
	wsURL := util.MakeWsURL(server.URL)
	defer server.Close()

	// mechanism to know test has completed
	done := make(chan bool, 1)

	clientHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !websocket.IsWebSocketUpgrade(r) {
			http.NotFound(w, r)
		}

		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			t.Fatal(err)
		}

		for {
			_, _, err = conn.NextReader()
			if err != nil && websocket.IsCloseError(err, websocket.CloseAbnormalClosure) {
				logger.Printf("closed")
				close(done)
				break
			}
			if err != nil {
				t.Fatal(err)
			}
		}

	})

	// register worker and serve http
	clientWs, _, err := websocket.DefaultDialer.Dial(wsURL+"/register/wsWorker/", nil)
	if err != nil {
		t.Fatal(err)
	}

	clientServer := &http.Server{Handler: clientHandler}
	go func() {
		_ = clientServer.Serve(wsmux.Client(clientWs, wsmux.Config{}))
	}()
	defer func() {
		_ = clientServer.Close()
	}()

	// create websocket connection
	conn, _, err := websocket.DefaultDialer.Dial(wsURL+"/wsWorker/", nil)
	if err != nil {
		t.Fatal(err)
	}

	// set timer for timing out test
	timer := time.NewTimer(4 * time.Second)

	// Close connection
	// will cause abnormal closure as Close will cause the underlying connection
	// to close without sending any close frame
	err = conn.Close()
	if err != nil {
		t.Fatal(err)
	}

	select {
	case <-timer.C:
		t.Fatalf("test failed: timeout")
	case <-done:
	}

}

// Ensures that session is removed once websocket connection is closed
func TestProxySessionRemoved(t *testing.T) {
	done := make(chan bool, 1)
	proxy := New(Config{Upgrader: upgrader, Logger: genLogger("proxy-session-remove-test")})
	proxy.SetSessionRemoveHandler(func(id string) {
		close(done)
	})

	server := httptest.NewServer(proxy.GetHandler())
	defer server.Close()

	wsURL := util.MakeWsURL(server.URL)
	conn, _, err := websocket.DefaultDialer.Dial(wsURL+"/register/wsWorker", nil)
	if err != nil {
		t.Fatal(err)
	}

	timeout := time.NewTimer(4 * time.Second)
	err = conn.Close()
	if err != nil {
		t.Fatal(err)
	}

	select {
	case <-done:
	case <-timeout.C:
		t.Fatalf("test timed out")
	}
}
<commit_msg>proxy: control test is now bidirectional<commit_after>package whproxy

import (
	"bytes"
	"io"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"
	"time"

	"sync"

	"github.com/gorilla/websocket"
	log "github.com/sirupsen/logrus"
	"github.com/taskcluster/webhooktunnel/util"
	"github.com/taskcluster/webhooktunnel/wsmux"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  64 * 1024,
	WriteBufferSize: 64 * 1024,
}

func genLogger(fname string) *log.Logger {
	file, err := os.Create(fname)
	if err != nil {
		panic(err)
	}
	logger := &log.Logger{
		Out:       file,
		Formatter: new(log.TextFormatter),
		Level:     log.DebugLevel,
	}
	return logger
}

func TestProxyRegister(t *testing.T) {
	//  start proxy server
	proxy := New(Config{Upgrader: upgrader, Logger: genLogger("register-test")})
	server := httptest.NewServer(proxy.GetHandler())
	defer server.Close()

	// get url
	wsURL := util.MakeWsURL(server.URL)

	// create address to dial
	workerID := "validWorkerID"
	dialAddr := wsURL + "/register/" + workerID

	// dial connection to proxy
	conn1, _, err := websocket.DefaultDialer.Dial(dialAddr, nil)
	if err != nil {
		t.Fatal(err)
	}

	defer func() {
		_ = conn1.Close()
	}()
	// second connection should fail
	conn2, _, err := websocket.DefaultDialer.Dial(dialAddr, nil)
	if err == nil {
		defer func() {
			_ = conn2.Close()
		}()
		t.Fatalf("bad status code: connection should fail")
	}
}

// TestProxyRequest
func TestProxyRequest(t *testing.T) {
	proxy := New(Config{Upgrader: upgrader, Logger: genLogger("request-test")})
	server := httptest.NewServer(proxy.GetHandler())
	defer server.Close()

	// get url
	wsURL := util.MakeWsURL(server.URL)
	// makeshift client
	clientWs, _, err := websocket.DefaultDialer.Dial(wsURL+"/register/workerID", nil)
	if err != nil {
		t.Fatal(err)
	}

	// handler to serve client requests
	clientHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.Method {
		case http.MethodGet:
			_, _ = w.Write([]byte("GET successful"))
		case http.MethodPost:
			_, _ = io.Copy(w, r.Body)
		default:
			http.NotFound(w, r)
		}
	})

	// serve client endpoint
	clientServer := &http.Server{Handler: clientHandler}
	go func() {
		_ = clientServer.Serve(wsmux.Client(clientWs, wsmux.Config{}))
	}()
	defer func() {
		_ = clientServer.Close()
	}()

	// make requests
	viewer := &http.Client{}
	servURL := server.URL

	// GET request
	resp, err := viewer.Get(servURL + "/workerID/")
	if err != nil {
		t.Fatal(err)
	}
	if resp.StatusCode != 200 {
		t.Log(resp)
		t.Fatalf("bad status code on get request")
	}
	reply, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(reply, []byte("GET successful")) {
		t.Fatalf("GET failed. Bad message")
	}

	// POST request
	resp, err = viewer.Post(servURL+"/workerID/", "application/text", bytes.NewBuffer([]byte("message")))
	if err != nil {
		t.Fatal(err)
	}
	if resp.StatusCode != 200 {
		t.Fatalf("bad status code on post request")
	}
	reply, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(reply, []byte("message")) {
		t.Fatalf("POST failed. Bad message")
	}

	// GET request to invalid id
	resp, err = viewer.Get(servURL + "/notWorkerID/")
	if resp.StatusCode != 404 {
		t.Fatalf("request should fail with 404")
	}
}

func TestProxyWebsocket(t *testing.T) {
	proxy := New(Config{Upgrader: upgrader})
	server := httptest.NewServer(proxy.GetHandler())
	wsURL := util.MakeWsURL(server.URL)
	defer server.Close()

	clientHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !websocket.IsWebSocketUpgrade(r) {
			http.NotFound(w, r)
		}

		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			t.Fatal(err)
		}

		mt, buf, err := conn.ReadMessage()
		if err != nil {
			t.Fatal(err)
		}

		err = conn.WriteMessage(mt, buf)
		if err != nil {
			t.Fatal(err)
		}
	})

	// register worker and serve http
	clientWs, _, err := websocket.DefaultDialer.Dial(wsURL+"/register/wsWorker/", nil)
	if err != nil {
		t.Fatal(err)
	}

	clientServer := &http.Server{Handler: clientHandler}
	go func() {
		_ = clientServer.Serve(wsmux.Client(clientWs, wsmux.Config{}))
	}()
	defer func() {
		_ = clientServer.Close()
	}()

	// create websocket connection
	conn, _, err := websocket.DefaultDialer.Dial(wsURL+"/wsWorker/", nil)
	if err != nil {
		t.Fatal(err)
	}
	defer func() {
		_ = conn.Close()
	}()

	// Generate 1M message
	message := make([]byte, 0)
	for i := 0; i < 1024*1024; i++ {
		message = append(message, byte(i%127))
	}

	err = conn.WriteMessage(websocket.BinaryMessage, message)
	if err != nil {
		t.Fatal(err)
	}

	_, buf, err := conn.ReadMessage()
	if !bytes.Equal(buf, message) {
		t.Fatalf("websocket test failed. Bad message")
	}
}

// ensure control messages are proxied
func TestWebsocketProxyControl(t *testing.T) {
	logger := genLogger("ws-control-test")
	proxy := New(Config{Upgrader: upgrader})
	//serve proxy
	server := httptest.NewServer(proxy.GetHandler())
	wsURL := util.MakeWsURL(server.URL)
	defer server.Close()

	// mechanism to know test has completed
	var wg sync.WaitGroup
	wg.Add(4)
	done := func() chan bool {
		tdone := make(chan bool, 1)
		go func() {
			wg.Wait()
			close(tdone)
		}()
		return tdone
	}

	clientHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !websocket.IsWebSocketUpgrade(r) {
			http.NotFound(w, r)
		}

		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			t.Fatal(err)
		}

		// set ping handler. Decrement wg to ensure ping frame was received
		conn.SetPingHandler(func(appData string) error {
			defer wg.Done()
			return conn.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(500*time.Millisecond))
		})

		// set pong handler. Decrement wg when called.
		conn.SetPongHandler(func(appData string) error {
			defer wg.Done()
			logger.Printf("received pong: %s", appData)
			if appData != "ping" {
				t.Fatal("bad pong")
			}
			return nil
		})

		err = conn.WriteControl(websocket.PingMessage, []byte("ping"), time.Now().Add(1*time.Second))
		// Read message to make sure ping was received
		for {
			_, _, err = conn.NextReader()
			if err != nil {
				break
			}
		}
	})

	// register worker and serve http
	clientWs, _, err := websocket.DefaultDialer.Dial(wsURL+"/register/wsWorker/", nil)
	if err != nil {
		t.Fatal(err)
	}

	clientServer := &http.Server{Handler: clientHandler}
	go func() {
		_ = clientServer.Serve(wsmux.Client(clientWs, wsmux.Config{}))
	}()
	defer func() {
		_ = clientServer.Close()
	}()

	// create websocket connection
	conn, _, err := websocket.DefaultDialer.Dial(wsURL+"/wsWorker/", nil)
	if err != nil {
		t.Fatal(err)
	}
	defer func() {
		_ = conn.Close()
	}()

	// Set Pong Handler. Decrement wg when pong handler fires to ensure that
	// pong is called
	conn.SetPongHandler(func(appData string) error {
		defer wg.Done()
		logger.Printf("received pong: %s", appData)
		if appData != "ping" {
			t.Fatal("bad pong")
		}
		return nil
	})

	conn.SetPingHandler(func(appData string) error {
		defer wg.Done()
		return conn.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(500*time.Millisecond))
	})

	// set timer for timing out test
	timer := time.NewTimer(3 * time.Second)

	// start reading messages to ensure pong is received
	go func() {
		for {
			_, _, err = conn.NextReader()
			if err != nil {
				break
			}
		}
	}()

	err = conn.WriteControl(websocket.PingMessage, []byte("ping"), time.Now().Add(1*time.Second))
	if err != nil {
		t.Fatal(err)
	}

	select {
	case <-timer.C:
		t.Fatalf("test failed: timeout")
	case <-done():
	}

}

// Ensure websocket close is proxied
func TestWebSocketClosure(t *testing.T) {
	logger := genLogger("ws-closure-test")
	proxy := New(Config{Upgrader: upgrader})
	//serve proxy
	server := httptest.NewServer(proxy.GetHandler())
	wsURL := util.MakeWsURL(server.URL)
	defer server.Close()

	// mechanism to know test has completed
	done := make(chan bool, 1)

	clientHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !websocket.IsWebSocketUpgrade(r) {
			http.NotFound(w, r)
		}

		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			t.Fatal(err)
		}

		for {
			_, _, err = conn.NextReader()
			if err != nil && websocket.IsCloseError(err, websocket.CloseAbnormalClosure) {
				logger.Printf("closed")
				close(done)
				break
			}
			if err != nil {
				t.Fatal(err)
			}
		}

	})

	// register worker and serve http
	clientWs, _, err := websocket.DefaultDialer.Dial(wsURL+"/register/wsWorker/", nil)
	if err != nil {
		t.Fatal(err)
	}

	clientServer := &http.Server{Handler: clientHandler}
	go func() {
		_ = clientServer.Serve(wsmux.Client(clientWs, wsmux.Config{}))
	}()
	defer func() {
		_ = clientServer.Close()
	}()

	// create websocket connection
	conn, _, err := websocket.DefaultDialer.Dial(wsURL+"/wsWorker/", nil)
	if err != nil {
		t.Fatal(err)
	}

	// set timer for timing out test
	timer := time.NewTimer(4 * time.Second)

	// Close connection
	// will cause abnormal closure as Close will cause the underlying connection
	// to close without sending any close frame
	err = conn.Close()
	if err != nil {
		t.Fatal(err)
	}

	select {
	case <-timer.C:
		t.Fatalf("test failed: timeout")
	case <-done:
	}

}

// Ensures that session is removed once websocket connection is closed
func TestProxySessionRemoved(t *testing.T) {
	done := make(chan bool, 1)
	proxy := New(Config{Upgrader: upgrader, Logger: genLogger("proxy-session-remove-test")})
	proxy.SetSessionRemoveHandler(func(id string) {
		close(done)
	})

	server := httptest.NewServer(proxy.GetHandler())
	defer server.Close()

	wsURL := util.MakeWsURL(server.URL)
	conn, _, err := websocket.DefaultDialer.Dial(wsURL+"/register/wsWorker", nil)
	if err != nil {
		t.Fatal(err)
	}

	timeout := time.NewTimer(4 * time.Second)
	err = conn.Close()
	if err != nil {
		t.Fatal(err)
	}

	select {
	case <-done:
	case <-timeout.C:
		t.Fatalf("test timed out")
	}
}
<|endoftext|>


<commit_before>package gokeepasslib

import (
	"bytes"
	"compress/gzip"
	"crypto/aes"
	"crypto/cipher"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
)

type Decoder struct {
	r io.Reader
}

func (d *Decoder) Decode(db *Database) error {
	if err := d.readSignature(db); err != nil {
		return err
	}

	if err := d.readHeaders(db); err != nil {
		return err
	}

	if err := d.readData(db); err != nil {
		return err
	}

	return nil
}

func NewDecoder(r io.Reader) *Decoder {
	return &Decoder{r: r}
}

func (d *Decoder) readSignature(db *Database) error {
	sig := new(Signature)
	if err := binary.Read(d.r, binary.LittleEndian, sig); err != nil {
		return err
	}

	if sig.BaseSignature != [...]byte{0x03, 0xd9, 0xa2, 0x9a} {
		return errors.New("BaseSignature not valid")
	}
	if sig.VersionSignature != [...]byte{0x67, 0xfb, 0x4b, 0xb5} {
		return errors.New("VersionSignature not valid")
	}

	db.signature = *sig
	return nil
}

func (d *Decoder) readHeaders(db *Database) error {
	headers := new(Headers)

	for {
		var fieldID byte
		if err := binary.Read(d.r, binary.LittleEndian, &fieldID); err != nil {
			return err
		}

		var fieldLength [2]byte
		if err := binary.Read(d.r, binary.LittleEndian, &fieldLength); err != nil {
			return err
		}

		var fieldData = make([]byte, binary.LittleEndian.Uint16(fieldLength[:]))
		if err := binary.Read(d.r, binary.LittleEndian, &fieldData); err != nil {
			return err
		}

		switch fieldID {
		case 1:
			headers.Comment = fieldData
		case 2:
			headers.CipherID = fieldData
		case 3:
			headers.CompressionFlags = binary.LittleEndian.Uint32(fieldData)
		case 4:
			headers.MasterSeed = fieldData
		case 5:
			headers.TransformSeed = fieldData
		case 6:
			data := binary.LittleEndian.Uint32(fieldData[4:8])*
				(1<<16)*(1<<16) +
				binary.LittleEndian.Uint32(fieldData[0:4])
			headers.TransformRounds = data
		case 7:
			headers.EncryptionIV = fieldData
		case 8:
			headers.ProtectedStreamKey = fieldData
		case 9:
			headers.StreamStartBytes = fieldData
		case 10:
			headers.InnerRandomStreamID = fieldData
		}

		if fieldID == 0 {
			break
		}
	}

	db.headers = *headers
	return nil
}

func (d *Decoder) readData(db *Database) error {
	block, err := aes.NewCipher(db.credentials.Key)
	if err != nil {
		return err
	}

	in, err := ioutil.ReadAll(d.r)
	if err != nil {
		return err
	}

	mode := cipher.NewCBCDecrypter(block, db.headers.EncryptionIV)
	decrypted := make([]byte, len(in))
	mode.CryptBlocks(decrypted, in)

	b := bytes.NewBuffer(decrypted)
	r, err := gzip.NewReader(b)
	if err != nil {
		return err
	}
	defer r.Close()
	result := []byte{}
	r.Read(result)

	fmt.Println(result)

	return nil
}
<commit_msg>trying to figure out aes ecb 256 in go...<commit_after>package gokeepasslib

import (
	"bytes"
	"compress/gzip"
	"crypto/aes"
	"crypto/cipher"
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
)

type Decoder struct {
	r io.Reader
}

func (d *Decoder) Decode(db *Database) error {
	if err := d.readSignature(db); err != nil {
		return err
	}

	if err := d.readHeaders(db); err != nil {
		return err
	}

	if err := d.readData(db); err != nil {
		return err
	}

	return nil
}

func NewDecoder(r io.Reader) *Decoder {
	return &Decoder{r: r}
}

func (d *Decoder) readSignature(db *Database) error {
	sig := new(Signature)
	if err := binary.Read(d.r, binary.LittleEndian, sig); err != nil {
		return err
	}

	if sig.BaseSignature != [...]byte{0x03, 0xd9, 0xa2, 0x9a} {
		return errors.New("BaseSignature not valid")
	}
	if sig.VersionSignature != [...]byte{0x67, 0xfb, 0x4b, 0xb5} {
		return errors.New("VersionSignature not valid")
	}

	db.signature = *sig
	return nil
}

func (d *Decoder) readHeaders(db *Database) error {
	headers := new(Headers)

	for {
		var fieldID byte
		if err := binary.Read(d.r, binary.LittleEndian, &fieldID); err != nil {
			return err
		}

		var fieldLength [2]byte
		if err := binary.Read(d.r, binary.LittleEndian, &fieldLength); err != nil {
			return err
		}

		var fieldData = make([]byte, binary.LittleEndian.Uint16(fieldLength[:]))
		if err := binary.Read(d.r, binary.LittleEndian, &fieldData); err != nil {
			return err
		}

		switch fieldID {
		case 1:
			headers.Comment = fieldData
		case 2:
			headers.CipherID = fieldData
		case 3:
			headers.CompressionFlags = binary.LittleEndian.Uint32(fieldData)
		case 4:
			headers.MasterSeed = fieldData
		case 5:
			headers.TransformSeed = fieldData
		case 6:
			data := binary.LittleEndian.Uint32(fieldData[4:8])*
				(1<<16)*(1<<16) +
				binary.LittleEndian.Uint32(fieldData[0:4])
			headers.TransformRounds = data
		case 7:
			headers.EncryptionIV = fieldData
		case 8:
			headers.ProtectedStreamKey = fieldData
		case 9:
			headers.StreamStartBytes = fieldData
		case 10:
			headers.InnerRandomStreamID = fieldData
		}

		if fieldID == 0 {
			break
		}
	}

	db.headers = *headers
	return nil
}

func (d *Decoder) readData(db *Database) error {
	masterKey, nil := d.buildMasterKey(db)

	block, err := aes.NewCipher(masterKey)
	if err != nil {
		return err
	}

	in, err := ioutil.ReadAll(d.r)
	if err != nil {
		return err
	}

	mode := cipher.NewCBCDecrypter(block, db.headers.EncryptionIV)
	decrypted := make([]byte, len(in))
	mode.CryptBlocks(decrypted, in)

	_ = db.headers.StreamStartBytes
	//fmt.Printf("%b\n", startBytes)
	//fmt.Printf("%b\n", decrypted[0:len(startBytes)])
	// if decrypted[0:len(startBytes)] != startBytes[:] {
	// 	return errors.New("Database integrity check failed")
	// }

	b := bytes.NewBuffer(decrypted)
	r, err := gzip.NewReader(b)
	if err != nil {
		return err
	}
	defer r.Close()
	result := []byte{}
	r.Read(result)

	fmt.Println(result)

	return nil
}

func (d *Decoder) buildMasterKey(db *Database) ([]byte, error) {
	masterKey := make([]byte, 32)
	copy(masterKey, db.credentials.Key)

	tmp := sha256.Sum256(masterKey)
	masterKey = tmp[:]

	block, err := aes.NewCipher(db.headers.TransformSeed)
	if err != nil {
		return nil, err
	}

	fmt.Printf("Before: % x\n", masterKey)

	//http://crypto.stackexchange.com/questions/21048/can-i-simulate-iterated-aes-ecb-with-other-block-cipher-modes

	for i := uint32(0); i < 1; i++ {
		result := make([]byte, 16)
		crypter := cipher.NewCBCEncrypter(block, result)
		crypter.CryptBlocks(result, masterKey[:16])
		// copy(masterKey[:16], result[:16])
		// result = make([]byte, 16)
		// crypter = cipher.NewCBCEncrypter(block, masterKey[16:])
		// crypter.CryptBlocks(result, result)
		// copy(masterKey[16:], result[16:])
		//ecbCrypt(result, masterKey, block)
	}

	fmt.Printf("After:  % x\n", masterKey)
	fmt.Printf("Should be: % x\n", []byte{0x20, 0x07, 0xbb, 0x4b, 0xdc, 0xa2, 0x86, 0xfb, 0x92, 0x50, 0xf3, 0x9f, 0x11, 0x1b, 0xbf, 0x77, 0x7e, 0x07, 0xd3, 0x80, 0x7c, 0x4a, 0x4e, 0x57, 0xbb, 0xc3, 0x89, 0x4f, 0x30, 0x4e, 0x4c, 0x1f})
	// 06 f4 38 80 ef 1c f2 fa c8 64 e4 4d 9c 42 eb 74 33 e9 22 07 c4 0a d3 29 0e 07 af c3 03 7a f9 5b

	masterKey = append(db.headers.MasterSeed, masterKey...)
	masterHash := sha256.Sum256(masterKey)
	masterKey = masterHash[:]

	return masterKey, nil
}

func ecbCrypt(dst, src []byte, block cipher.Block) {
	length := len(src) / block.BlockSize()
	for i := 0; i < length; i++ {
		block.Encrypt(dst, src)
		src = src[block.BlockSize():]
		dst = dst[block.BlockSize():]
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 The Ebiten 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 ebitenutil

import (
	"image/color"
	"math"

	"github.com/hajimehoshi/ebiten"
)

var (
	emptyNearestImage *ebiten.Image
	emptyLinearImage  *ebiten.Image
)

func init() {
	emptyNearestImage, _ = ebiten.NewImage(16, 16, ebiten.FilterNearest)
	_ = emptyNearestImage.Fill(color.White)
	emptyLinearImage, _ = ebiten.NewImage(16, 16, ebiten.FilterLinear)
	_ = emptyLinearImage.Fill(color.White)
}

func colorScale(clr color.Color) (rf, gf, bf, af float64) {
	r, g, b, a := clr.RGBA()
	if a == 0 {
		return 0, 0, 0, 0
	}

	rf = float64(r) / float64(a)
	gf = float64(g) / float64(a)
	bf = float64(b) / float64(a)
	af = float64(a) / 0xffff
	return
}

// DrawLine draws a line segment on the given destination dst.
//
// DrawLine is intended to be used mainly for debugging or prototyping purpose.
func DrawLine(dst *ebiten.Image, x1, y1, x2, y2 float64, clr color.Color) {
	ew, eh := emptyLinearImage.Size()
	length := math.Hypot(x2-x1, y2-y1)

	op := &ebiten.DrawImageOptions{}
	op.GeoM.Scale(length/float64(ew), 1/float64(eh))
	op.GeoM.Rotate(math.Atan2(y2-y1, x2-x1))
	op.GeoM.Translate(x1, y1)
	op.ColorM.Scale(colorScale(clr))
	_ = dst.DrawImage(emptyLinearImage, op)
}

// DrawRect draws a rectangle on the given destination dst.
//
// DrawRect is intended to be used mainly for debugging or prototyping purpose.
func DrawRect(dst *ebiten.Image, x, y, width, height float64, clr color.Color) {
	ew, eh := emptyNearestImage.Size()

	op := &ebiten.DrawImageOptions{}
	op.GeoM.Scale(width/float64(ew), height/float64(eh))
	op.GeoM.Translate(x, y)
	op.ColorM.Scale(colorScale(clr))
	_ = dst.DrawImage(emptyNearestImage, op)
}
<commit_msg>ebitenutil: Give up using linear filter for lines<commit_after>// Copyright 2017 The Ebiten 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 ebitenutil

import (
	"image/color"
	"math"

	"github.com/hajimehoshi/ebiten"
)

var (
	emptyImage *ebiten.Image
)

func init() {
	// Filter must be 'nearest' filter.
	// Linear filtering would make edges blurred.
	emptyImage, _ = ebiten.NewImage(16, 16, ebiten.FilterNearest)
	_ = emptyImage.Fill(color.White)
}

func colorScale(clr color.Color) (rf, gf, bf, af float64) {
	r, g, b, a := clr.RGBA()
	if a == 0 {
		return 0, 0, 0, 0
	}

	rf = float64(r) / float64(a)
	gf = float64(g) / float64(a)
	bf = float64(b) / float64(a)
	af = float64(a) / 0xffff
	return
}

// DrawLine draws a line segment on the given destination dst.
//
// DrawLine is intended to be used mainly for debugging or prototyping purpose.
func DrawLine(dst *ebiten.Image, x1, y1, x2, y2 float64, clr color.Color) {
	ew, eh := emptyImage.Size()
	length := math.Hypot(x2-x1, y2-y1)

	op := &ebiten.DrawImageOptions{}
	op.GeoM.Scale(length/float64(ew), 1/float64(eh))
	op.GeoM.Rotate(math.Atan2(y2-y1, x2-x1))
	op.GeoM.Translate(x1, y1)
	op.ColorM.Scale(colorScale(clr))
	_ = dst.DrawImage(emptyImage, op)
}

// DrawRect draws a rectangle on the given destination dst.
//
// DrawRect is intended to be used mainly for debugging or prototyping purpose.
func DrawRect(dst *ebiten.Image, x, y, width, height float64, clr color.Color) {
	ew, eh := emptyImage.Size()

	op := &ebiten.DrawImageOptions{}
	op.GeoM.Scale(width/float64(ew), height/float64(eh))
	op.GeoM.Translate(x, y)
	op.ColorM.Scale(colorScale(clr))
	_ = dst.DrawImage(emptyImage, op)
}
<|endoftext|>


<commit_before>package plugin

import (
	"zvr/server"
	"zvr/utils"
	"fmt"
	"strings"
	log "github.com/Sirupsen/logrus"
)

const (
	VR_CREATE_VIP = "/createvip"
	VR_REMOVE_VIP = "/removevip"
	VR_SET_VIP_QOS = "/setvipqos"
	VR_DELETE_VIP_QOS = "/deletevipqos"
	VR_SYNC_VIP_QOS = "/syncvipqos"
	VR_IFB = "ifb0"
)

type vipInfo struct {
	Ip string `json:"ip"`
	Netmask string `json:"netmask"`
	Gateway string `json:"gateway"`
	OwnerEthernetMac string `json:"ownerEthernetMac"`
}

type vipQosSettings struct {
	Ip                          string `json:"ip"`
	QosClassId                  int `json:"qosClassId"`
	VipUuid                     string `json:"vipUuid"`
	IsFirstSetInboundBandwidth  bool `json:"isFirstSetInboundBandwidth"`
	IsFirstSetOutboundBandwidth bool `json:"isFirstSetOutboundBandwidth"`
	InboundBandwidth            int64 `json:"inboundBandwidth"`
	OutBoundBandwidth           int64 `json:"outboundBandwidth"`
	NicMac						string `json:"nicMac"`
}

type setVipCmd struct {
	Vips []vipInfo `json:"vips"`
}

type removeVipCmd struct {
	Vips []vipInfo `json:"vips"`
}

type setVipQosCmd struct {
	Settings []vipQosSettings `json:"vipQosSettings"`

}

type deleteVipQosCmd struct {
	Settings []vipQosSettings `json:"vipQosSettings"`
}

type syncVipQosCmd struct {
	Settings []vipQosSettings `json:"vipQosSettings"`
}

func setVip(ctx *server.CommandContext) interface{} {
	cmd := &setVipCmd{}
	ctx.GetCommand(cmd)

	tree := server.NewParserFromShowConfiguration().Tree
	for _, vip := range cmd.Vips {
		nicname, err := utils.GetNicNameByMac(vip.OwnerEthernetMac); utils.PanicOnError(err)
		cidr, err := utils.NetmaskToCIDR(vip.Netmask); utils.PanicOnError(err)
		addr := fmt.Sprintf("%v/%v", vip.Ip, cidr)
		if n := tree.Getf("interfaces ethernet %s address %v", nicname, addr); n == nil {
			tree.SetfWithoutCheckExisting("interfaces ethernet %s address %v", nicname, addr)
		}
	}

	tree.Apply(false)

	return nil
}

func removeVip(ctx *server.CommandContext) interface{} {
	cmd := &removeVipCmd{}
	ctx.GetCommand(cmd)

	tree := server.NewParserFromShowConfiguration().Tree
	for _, vip := range cmd.Vips {
		nicname, err := utils.GetNicNameByMac(vip.OwnerEthernetMac); utils.PanicOnError(err)
		cidr, err := utils.NetmaskToCIDR(vip.Netmask); utils.PanicOnError(err)
		addr := fmt.Sprintf("%v/%v", vip.Ip, cidr)

		tree.Deletef("interfaces ethernet %s address %v", nicname, addr)
	}

	tree.Apply(false)

	return nil
}

func setVipQos(ctx *server.CommandContext) interface{} {
	cmd := &setVipQosCmd{}
	ctx.GetCommand(cmd)

	for _, setting := range cmd.Settings {
		addr := fmt.Sprintf("%v/32", setting.Ip)
		publicInterface, err := utils.GetNicNameByIp(setting.Ip); utils.PanicOnError(err)
		if(setting.OutBoundBandwidth != -1 && setting.OutBoundBandwidth != 0){
			initPublicInterfaceQosConfigurationIfNotInit(publicInterface)
			addClass(publicInterface, setting.OutBoundBandwidth, setting.QosClassId)
			if(setting.IsFirstSetOutboundBandwidth){
				addOutFilter(publicInterface, addr, setting.QosClassId)
			}
		}

		if(setting.InboundBandwidth != -1 && setting.InboundBandwidth != 0){
			initIfbConfigurationIfNotInit(publicInterface)
			addClass(VR_IFB, setting.InboundBandwidth, setting.QosClassId)
			if(setting.IsFirstSetInboundBandwidth){
				addInFilter(VR_IFB, addr, setting.QosClassId)
			}
		}

	}

	return nil
}

func deleteVipQos(ctx *server.CommandContext) interface{} {
	cmd := &deleteVipQosCmd{}
	ctx.GetCommand(cmd)

	for _, setting := range cmd.Settings {
		publicInterface, error := utils.GetNicNameByIp(setting.Ip)
		if (error != nil && setting.NicMac != "") {
			log.Debugf("no nic for ip %s found, try to use mac %s", setting.Ip, setting.NicMac)
			publicInterface, error = utils.GetNicNameByMac(setting.NicMac); utils.PanicOnError(error)
		}
		if(setting.OutBoundBandwidth == 0){
			delFilter(publicInterface, setting.QosClassId)
			delClass(publicInterface, setting.QosClassId)
		}
		if(setting.InboundBandwidth == 0){
			delFilter(VR_IFB, setting.QosClassId)
			delClass(VR_IFB, setting.QosClassId)
		}
	}

	return nil
}

func syncVipQos(ctx *server.CommandContext) interface{} {
	cmd := &syncVipQosCmd{}
	ctx.GetCommand(cmd)

	for _, setting := range cmd.Settings {
		addr := fmt.Sprintf("%v/32", setting.Ip)
		publicInterface, err := utils.GetNicNameByIp(setting.Ip); utils.PanicOnError(err)
		initPublicInterfaceQosConfigurationIfNotInit(publicInterface)
		initIfbConfigurationIfNotInit(publicInterface)
		if(setting.OutBoundBandwidth != -1 && setting.OutBoundBandwidth != 0){
			addClass(publicInterface, setting.OutBoundBandwidth, setting.QosClassId)
			addOutFilter(publicInterface, addr, setting.QosClassId)
		}
		if(setting.InboundBandwidth != -1 && setting.InboundBandwidth != 0){
			addClass(VR_IFB, setting.InboundBandwidth, setting.QosClassId)
			addInFilter(VR_IFB, addr, setting.QosClassId)
		}

	}

	return nil
}

func addClass(nicName string, bandwidth int64, qosClassId int){
	bash := utils.Bash{
		Command: fmt.Sprintf("sudo tc class replace dev %s parent 1: classid 1:%d htb rate %dbit ceil %dbit burst 10k cburst 10k", nicName, qosClassId, bandwidth, bandwidth),
	}
	bash.Run()
	bash.PanicIfError()
}

func delClass(nicName string, qosClassId int){
	bash := utils.Bash{
		Command: fmt.Sprintf("sudo tc class del dev %s parent 1: classid 1:%d", nicName, qosClassId),
	}
	bash.Run()
	bash.PanicIfError()
}

func addOutFilter(nicName string, addr string, qosClassId int){
	bash := utils.Bash{
		Command:fmt.Sprintf("sudo tc filter replace dev %s parent 1: protocol ip prio 12 u32 match ip src %s flowid 1:%d", nicName, addr, qosClassId),
	}
	bash.Run()
	bash.PanicIfError()
}

func addInFilter(nicName string, addr string, qosClassId int){
	bash := utils.Bash{
		Command:fmt.Sprintf("sudo tc filter replace dev %s parent 1: protocol ip prio 12 u32 match ip dst %s flowid 1:%d", nicName, addr, qosClassId),
	}
	bash.Run()
	bash.PanicIfError()
}

func delFilter(nicName string, qosClassId int){
	bash := utils.Bash{
		Command:fmt.Sprintf("handles=`sudo tc filter list dev %s |grep \"flowid 1:%d\" |awk '{print $10}'`; for handle in $handles; do sudo tc filter delete dev %s parent 1: protocol ip prio 12 handle ${handle} u32; done", nicName, qosClassId, nicName),
	}
	bash.Run()
	bash.PanicIfError()
}

func initIfbConfigurationIfNotInit(pubInterface string){
	//init ifb
	bash :=utils.Bash{
		Command: "ip a |grep ifb | awk '{print $2}' |sed 's/\\://g'",
	}
	_,o,_,_ := bash.RunWithReturn();bash.PanicIfError()
	o = strings.TrimSpace(o)
	if(!strings.Contains(o, VR_IFB)){
		bash = utils.Bash{
			Command: fmt.Sprintf("/sbin/modprobe ifb numifbs=0;" +
				"ip link add %s type ifb;" +
				"ip link set dev %s up;", VR_IFB, VR_IFB,),
		}

		bash.Run()
		bash.PanicIfError()
	}

	//init public interface ingress
	bash = utils.Bash{
		Command: fmt.Sprintf("sudo tc qdisc show dev %s | awk '{print $2,$3}'", pubInterface),
	}
	_,o,_,_ = bash.RunWithReturn();bash.PanicIfError()
	if(!strings.Contains(o, "ingress ffff:")){
		bash = utils.Bash{
			Command: fmt.Sprintf("sudo tc qdisc add dev %s handle ffff: ingress",pubInterface),
		}
		bash.Run()
		bash.PanicIfError()
	}

	//init ifb qos configuration
	bash = utils.Bash{
		Command: fmt.Sprintf("sudo tc qdisc show dev %s| awk '{print $2,$3,$4}'", VR_IFB),
	}
	_,o,_,_ = bash.RunWithReturn();bash.PanicIfError()
	o = strings.TrimSpace(o)
	if(!strings.Contains(o,"htb 1: root")){
		bash = utils.Bash{
			Command: fmt.Sprintf("sudo tc qdisc add dev %s root handle 1:0 htb default a;" +
				"sudo tc filter add dev %s parent ffff: protocol ip u32 match u32 0 0 action mirred egress redirect dev %s", VR_IFB, pubInterface, VR_IFB),
		}
		bash.Run()
		bash.PanicIfError()
	}
}

func initPublicInterfaceQosConfigurationIfNotInit(pubInterface string){
	bash := utils.Bash{
		Command: fmt.Sprintf("sudo tc qdisc show dev %s | awk '{print $2,$3,$4}'", pubInterface),
	}
	_,o,_,_ := bash.RunWithReturn(); bash.PanicIfError()
	o = strings.TrimSpace(o)
	if(!strings.Contains(o,"htb 1: root")){
		bash = utils.Bash{
			Command: fmt.Sprintf("sudo tc qdisc add dev %s root handle 1:0 htb default a", pubInterface),
		}
		bash.Run()
		bash.PanicIfError()
	}

}

func VipEntryPoint()  {
	server.RegisterAsyncCommandHandler(VR_CREATE_VIP, server.VyosLock(setVip))
	server.RegisterAsyncCommandHandler(VR_REMOVE_VIP, server.VyosLock(removeVip))
	server.RegisterAsyncCommandHandler(VR_SET_VIP_QOS, server.VyosLock(setVipQos))
	server.RegisterAsyncCommandHandler(VR_DELETE_VIP_QOS, server.VyosLock(deleteVipQos))
	server.RegisterAsyncCommandHandler(VR_SYNC_VIP_QOS, server.VyosLock(syncVipQos))
}
<commit_msg>Fix vip qos can will not work if set qos before use the vip<commit_after>package plugin

import (
	"zvr/server"
	"zvr/utils"
	"fmt"
	"strings"
	log "github.com/Sirupsen/logrus"
)

const (
	VR_CREATE_VIP = "/createvip"
	VR_REMOVE_VIP = "/removevip"
	VR_SET_VIP_QOS = "/setvipqos"
	VR_DELETE_VIP_QOS = "/deletevipqos"
	VR_SYNC_VIP_QOS = "/syncvipqos"
	VR_IFB = "ifb0"
)

type vipInfo struct {
	Ip string `json:"ip"`
	Netmask string `json:"netmask"`
	Gateway string `json:"gateway"`
	OwnerEthernetMac string `json:"ownerEthernetMac"`
}

type vipQosSettings struct {
	Ip                          string `json:"ip"`
	QosClassId                  int `json:"qosClassId"`
	VipUuid                     string `json:"vipUuid"`
	IsFirstSetInboundBandwidth  bool `json:"isFirstSetInboundBandwidth"`
	IsFirstSetOutboundBandwidth bool `json:"isFirstSetOutboundBandwidth"`
	InboundBandwidth            int64 `json:"inboundBandwidth"`
	OutBoundBandwidth           int64 `json:"outboundBandwidth"`
	NicMac						string `json:"nicMac"`
}

type setVipCmd struct {
	Vips []vipInfo `json:"vips"`
}

type removeVipCmd struct {
	Vips []vipInfo `json:"vips"`
}

type setVipQosCmd struct {
	Settings []vipQosSettings `json:"vipQosSettings"`

}

type deleteVipQosCmd struct {
	Settings []vipQosSettings `json:"vipQosSettings"`
}

type syncVipQosCmd struct {
	Settings []vipQosSettings `json:"vipQosSettings"`
}

func setVip(ctx *server.CommandContext) interface{} {
	cmd := &setVipCmd{}
	ctx.GetCommand(cmd)

	tree := server.NewParserFromShowConfiguration().Tree
	for _, vip := range cmd.Vips {
		nicname, err := utils.GetNicNameByMac(vip.OwnerEthernetMac); utils.PanicOnError(err)
		cidr, err := utils.NetmaskToCIDR(vip.Netmask); utils.PanicOnError(err)
		addr := fmt.Sprintf("%v/%v", vip.Ip, cidr)
		if n := tree.Getf("interfaces ethernet %s address %v", nicname, addr); n == nil {
			tree.SetfWithoutCheckExisting("interfaces ethernet %s address %v", nicname, addr)
		}
	}

	tree.Apply(false)

	return nil
}

func removeVip(ctx *server.CommandContext) interface{} {
	cmd := &removeVipCmd{}
	ctx.GetCommand(cmd)

	tree := server.NewParserFromShowConfiguration().Tree
	for _, vip := range cmd.Vips {
		nicname, err := utils.GetNicNameByMac(vip.OwnerEthernetMac); utils.PanicOnError(err)
		cidr, err := utils.NetmaskToCIDR(vip.Netmask); utils.PanicOnError(err)
		addr := fmt.Sprintf("%v/%v", vip.Ip, cidr)

		tree.Deletef("interfaces ethernet %s address %v", nicname, addr)
	}

	tree.Apply(false)

	return nil
}

func setVipQos(ctx *server.CommandContext) interface{} {
	cmd := &setVipQosCmd{}
	ctx.GetCommand(cmd)

	for _, setting := range cmd.Settings {
		addr := fmt.Sprintf("%v/32", setting.Ip)
		publicInterface, err := utils.GetNicNameByIp(setting.Ip); utils.PanicOnError(err)
		if(setting.OutBoundBandwidth != -1 && setting.OutBoundBandwidth != 0){
			initPublicInterfaceQosConfigurationIfNotInit(publicInterface)
			addClass(publicInterface, setting.OutBoundBandwidth, setting.QosClassId)
			if(setting.IsFirstSetOutboundBandwidth){
				addOutFilter(publicInterface, addr, setting.QosClassId)
			}
		}

		if(setting.InboundBandwidth != -1 && setting.InboundBandwidth != 0){
			initIfbConfigurationIfNotInit(publicInterface)
			addClass(VR_IFB, setting.InboundBandwidth, setting.QosClassId)
			if(setting.IsFirstSetInboundBandwidth){
				addInFilter(VR_IFB, addr, setting.QosClassId)
			}
		}

	}

	return nil
}

func deleteVipQos(ctx *server.CommandContext) interface{} {
	cmd := &deleteVipQosCmd{}
	ctx.GetCommand(cmd)

	for _, setting := range cmd.Settings {
		publicInterface, error := utils.GetNicNameByIp(setting.Ip)
		if (error != nil && setting.NicMac != "") {
			log.Debugf("no nic for ip %s found, try to use mac %s", setting.Ip, setting.NicMac)
			publicInterface, error = utils.GetNicNameByMac(setting.NicMac); utils.PanicOnError(error)
		}
		if(setting.OutBoundBandwidth == 0){
			delFilter(publicInterface, setting.QosClassId)
			delClass(publicInterface, setting.QosClassId)
		}
		if(setting.InboundBandwidth == 0){
			delFilter(VR_IFB, setting.QosClassId)
			delClass(VR_IFB, setting.QosClassId)
		}
	}

	return nil
}

func syncVipQos(ctx *server.CommandContext) interface{} {
	cmd := &syncVipQosCmd{}
	ctx.GetCommand(cmd)

	for _, setting := range cmd.Settings {
		addr := fmt.Sprintf("%v/32", setting.Ip)
		publicInterface, err := utils.GetNicNameByIp(setting.Ip);
		if (err != nil && setting.NicMac != "") {
			log.Debugf("no nic for ip %s found, try to use mac %s", setting.Ip, setting.NicMac)
			publicInterface, err = utils.GetNicNameByMac(setting.NicMac); utils.PanicOnError(err)
		}
		initPublicInterfaceQosConfigurationIfNotInit(publicInterface)
		initIfbConfigurationIfNotInit(publicInterface)
		if(setting.OutBoundBandwidth != -1 && setting.OutBoundBandwidth != 0){
			addClass(publicInterface, setting.OutBoundBandwidth, setting.QosClassId)
			addOutFilter(publicInterface, addr, setting.QosClassId)
		}
		if(setting.InboundBandwidth != -1 && setting.InboundBandwidth != 0){
			addClass(VR_IFB, setting.InboundBandwidth, setting.QosClassId)
			addInFilter(VR_IFB, addr, setting.QosClassId)
		}
		if(setting.OutBoundBandwidth == 0){
			delFilter(publicInterface, setting.QosClassId)
			delClass(publicInterface, setting.QosClassId)
		}
		if(setting.InboundBandwidth == 0){
			delFilter(VR_IFB, setting.QosClassId)
			delClass(VR_IFB, setting.QosClassId)
		}
	}

	return nil
}

func addClass(nicName string, bandwidth int64, qosClassId int){
	bash := utils.Bash{
		Command: fmt.Sprintf("sudo tc class replace dev %s parent 1: classid 1:%d htb rate %dbit ceil %dbit burst 10k cburst 10k", nicName, qosClassId, bandwidth, bandwidth),
	}
	bash.Run()
	bash.PanicIfError()
}

func delClass(nicName string, qosClassId int){
	bash := utils.Bash{
		Command: fmt.Sprintf("sudo tc class del dev %s parent 1: classid 1:%d", nicName, qosClassId),
	}
	bash.Run()
	bash.PanicIfError()
}

func addOutFilter(nicName string, addr string, qosClassId int){
	bash := utils.Bash{
		Command:fmt.Sprintf("sudo tc filter replace dev %s parent 1: protocol ip prio 12 u32 match ip src %s flowid 1:%d", nicName, addr, qosClassId),
	}
	bash.Run()
	bash.PanicIfError()
}

func addInFilter(nicName string, addr string, qosClassId int){
	bash := utils.Bash{
		Command:fmt.Sprintf("sudo tc filter replace dev %s parent 1: protocol ip prio 12 u32 match ip dst %s flowid 1:%d", nicName, addr, qosClassId),
	}
	bash.Run()
	bash.PanicIfError()
}

func delFilter(nicName string, qosClassId int){
	bash := utils.Bash{
		Command:fmt.Sprintf("handles=`sudo tc filter list dev %s |grep \"flowid 1:%d\" |awk '{print $10}'`; for handle in $handles; do sudo tc filter delete dev %s parent 1: protocol ip prio 12 handle ${handle} u32; done", nicName, qosClassId, nicName),
	}
	bash.Run()
	bash.PanicIfError()
}

func initIfbConfigurationIfNotInit(pubInterface string){
	//init ifb
	bash :=utils.Bash{
		Command: "ip a |grep ifb | awk '{print $2}' |sed 's/\\://g'",
	}
	_,o,_,_ := bash.RunWithReturn();bash.PanicIfError()
	o = strings.TrimSpace(o)
	if(!strings.Contains(o, VR_IFB)){
		bash = utils.Bash{
			Command: fmt.Sprintf("/sbin/modprobe ifb numifbs=0;" +
				"ip link add %s type ifb;" +
				"ip link set dev %s up;", VR_IFB, VR_IFB,),
		}

		bash.Run()
		bash.PanicIfError()
	}

	//init public interface ingress
	bash = utils.Bash{
		Command: fmt.Sprintf("sudo tc qdisc show dev %s | awk '{print $2,$3}'", pubInterface),
	}
	_,o,_,_ = bash.RunWithReturn();bash.PanicIfError()
	if(!strings.Contains(o, "ingress ffff:")){
		bash = utils.Bash{
			Command: fmt.Sprintf("sudo tc qdisc add dev %s handle ffff: ingress",pubInterface),
		}
		bash.Run()
		bash.PanicIfError()
	}

	//init ifb qos configuration
	bash = utils.Bash{
		Command: fmt.Sprintf("sudo tc qdisc show dev %s| awk '{print $2,$3,$4}'", VR_IFB),
	}
	_,o,_,_ = bash.RunWithReturn();bash.PanicIfError()
	o = strings.TrimSpace(o)
	if(!strings.Contains(o,"htb 1: root")){
		bash = utils.Bash{
			Command: fmt.Sprintf("sudo tc qdisc add dev %s root handle 1:0 htb default a;" +
				"sudo tc filter add dev %s parent ffff: protocol ip u32 match u32 0 0 action mirred egress redirect dev %s", VR_IFB, pubInterface, VR_IFB),
		}
		bash.Run()
		bash.PanicIfError()
	}
}

func initPublicInterfaceQosConfigurationIfNotInit(pubInterface string){
	bash := utils.Bash{
		Command: fmt.Sprintf("sudo tc qdisc show dev %s | awk '{print $2,$3,$4}'", pubInterface),
	}
	_,o,_,_ := bash.RunWithReturn(); bash.PanicIfError()
	o = strings.TrimSpace(o)
	if(!strings.Contains(o,"htb 1: root")){
		bash = utils.Bash{
			Command: fmt.Sprintf("sudo tc qdisc add dev %s root handle 1:0 htb default a", pubInterface),
		}
		bash.Run()
		bash.PanicIfError()
	}

}

func VipEntryPoint()  {
	server.RegisterAsyncCommandHandler(VR_CREATE_VIP, server.VyosLock(setVip))
	server.RegisterAsyncCommandHandler(VR_REMOVE_VIP, server.VyosLock(removeVip))
	server.RegisterAsyncCommandHandler(VR_SET_VIP_QOS, server.VyosLock(setVipQos))
	server.RegisterAsyncCommandHandler(VR_DELETE_VIP_QOS, server.VyosLock(deleteVipQos))
	server.RegisterAsyncCommandHandler(VR_SYNC_VIP_QOS, server.VyosLock(syncVipQos))
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"os/user"
	"strings"
	"sync"

	"github.com/getcfs/fuse"
	pb "github.com/getcfs/megacfs/formic/proto"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
)

type server struct {
	fs *fs
	wg sync.WaitGroup
}

func newserver(fs *fs) *server {
	s := &server{
		fs: fs,
	}
	return s
}

func (s *server) serve() error {
	defer s.wg.Wait()

	for {
		req, err := s.fs.conn.ReadRequest()
		if err != nil {
			if err == io.EOF {
				break
			}
			return err
		}
		s.wg.Add(1)
		go func() {
			defer s.wg.Done()
			s.fs.handle(req)
		}()
	}
	return nil
}

func debuglog(msg interface{}) {
	fmt.Fprintf(os.Stderr, "%v\n", msg)
}

type rpc struct {
	conn *grpc.ClientConn
	api  pb.ApiClient
}

func newrpc(conn *grpc.ClientConn) *rpc {
	r := &rpc{
		conn: conn,
		api:  pb.NewApiClient(conn),
	}

	return r
}

// NullWriter ...
type NullWriter int

func (NullWriter) Write([]byte) (int, error) { return 0, nil }

var cfsVersion string
var buildDate string
var commitVersion string
var goVersion string

type authResponse struct {
	Access struct {
		Token struct {
			ID string `json:"id"`
		} `json:"token"`
	} `json:"access"`
}

func auth(username string, apikey string) string {
	body := fmt.Sprintf(`{"auth":{"RAX-KSKEY:apiKeyCredentials":{"username": "%s","apiKey": "%s"}}}`, username, apikey)
	rbody := strings.NewReader(body)
	req, err := http.NewRequest("POST", "https://identity.api.rackspacecloud.com/v2.0/tokens", rbody)
	if err != nil {
		fmt.Printf("%v", err)
		os.Exit(1)
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		fmt.Printf("%v", err)
		os.Exit(1)
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		fmt.Println(resp.Status)
		os.Exit(1)
	}

	// parse token from response
	var authResp authResponse
	r, _ := ioutil.ReadAll(resp.Body)
	json.Unmarshal(r, &authResp)
	token := authResp.Access.Token.ID

	return token
}

func mount() error {
	f := flag.NewFlagSet("mount", flag.ContinueOnError)
	f.Usage = func() {
		fmt.Println("Usage:")
		fmt.Println("    cfs mount [-o option,...] <region>:<fsid> <mountpoint>")
		fmt.Println("Options:")
		fmt.Println("    -o debug          enables debug output")
		fmt.Println("    -o ro             mount the filesystem read only")
		fmt.Println("    -o allow_other    allow access to other users")
		fmt.Println("Examples:")
		fmt.Println("    cfs mount iad:11111111-1111-1111-1111-111111111111 /mnt/test")
		fmt.Println("    cfs mount -o debug,ro iad:11111111-1111-1111-1111-111111111111 /mnt/test")
		os.Exit(1)
	}
	var options string
	f.StringVar(&options, "o", "", "")
	f.Parse(flag.Args()[1:])
	if f.NArg() != 2 {
		f.Usage()
	}
	regionFsid := f.Args()[0]
	parts := strings.Split(regionFsid, ":")
	if len(parts) != 2 {
		fmt.Println("Invalid filesystem:", regionFsid)
		f.Usage()
	}
	region := strings.ToLower(parts[0])
	fsid := parts[1]
	mountpoint := f.Args()[1]
	// Verify mountpoint exists
	_, err := os.Stat(mountpoint)
	if os.IsNotExist(err) {
		fmt.Printf("Mount point %s does not exist\n", mountpoint)
		os.Exit(1)
	}
	// Verify region
	addr, ok := regions[region]
	if !ok {
		fmt.Println("Invalid region:", region)
		os.Exit(1)
	}

	// Setup grpc
	var opts []grpc.DialOption
	creds := credentials.NewTLS(&tls.Config{
		InsecureSkipVerify: true,
	})
	opts = append(opts, grpc.WithTransportCredentials(creds))
	conn, err := grpc.Dial(addr, opts...)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	defer conn.Close()

	// handle fuse mount options
	mountOptions := []fuse.MountOption{
		fuse.FSName("cfs"),
		fuse.Subtype("cfs"),
		fuse.DefaultPermissions(),
		fuse.MaxReadahead(128 * 1024),
		fuse.AsyncRead(),
		//fuse.WritebackCache(),
		fuse.AutoInvalData(),
	}

	// parse mount options string
	clargs := getArgs(options)

	// handle debug mount option
	_, debug := clargs["debug"]
	if !debug {
		log.SetFlags(0)
		log.SetOutput(ioutil.Discard)
	}

	// handle allow_other mount option
	_, allowOther := clargs["allow_other"]
	if allowOther {
		mountOptions = append(mountOptions, fuse.AllowOther())
	}

	// handle ro mount option
	_, readOnly := clargs["ro"]
	if readOnly {
		mountOptions = append(mountOptions, fuse.ReadOnly())
	}

	// perform fuse mount
	fusermountPath()
	cfs, err := fuse.Mount(mountpoint, mountOptions...)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	defer cfs.Close()

	// setup rpc client
	rpc := newrpc(conn)
	fs := newfs(cfs, rpc, fsid)
	err = fs.InitFs()
	if err != nil {
		fmt.Println(err)
		exec.Command("fusermount", "-uz", mountpoint).Run()
		os.Exit(1)
	}
	srv := newserver(fs)

	if err := srv.serve(); err != nil {
		fmt.Println(err)
		exec.Command("fusermount", "-uz", mountpoint).Run()
		os.Exit(1)
	}

	<-cfs.Ready
	if err := cfs.MountError; err != nil {
		fmt.Println(err)
		exec.Command("fusermount", "-uz", mountpoint).Run()
		os.Exit(1)
	}
	return nil
}

func main() {

	// try to read config file
	configured := false
	config := map[string]string{}
	user, err := user.Current()
	if err != nil {
		fmt.Printf("Unable to identify curent user: %v\n", err)
		os.Exit(1)
	}
	var configfile string
	if user.HomeDir != "" {
		configfile = user.HomeDir + "/.cfs.json"
	} else {
		configfile = ".cfs.json"
	}
	f, err := ioutil.ReadFile(configfile)
	if err == nil {
		json.Unmarshal(f, &config)
		configured = true
	}
	region, _ := config["region"]
	username, _ := config["username"]
	apikey, _ := config["apikey"]

	flag.Usage = func() {
		fmt.Println("Usage:")
		fmt.Println("    cfs <command> [options] <args>")
		fmt.Println("Commands:")
		fmt.Println("    configure    create configuration file")
		fmt.Println("    list         list all filesystems")
		fmt.Println("    create       create a new filesytem")
		fmt.Println("    show         show filesystem details")
		fmt.Println("    update       update filesystem name")
		fmt.Println("    delete       delete an existing filesystem")
		fmt.Println("    grant        grant access to a filesystem")
		fmt.Println("    revoke       revoke access to a filesystem")
		fmt.Println("    mount        mount an existing filesystem")
		fmt.Println("    version      show client version")
		fmt.Println("    help         show usage for cfs")
		fmt.Println("Examples:")
		fmt.Println("    cfs configure")
		fmt.Println("    cfs create <name>")
		fmt.Println("    cfs grant <ip> <fsid>")
		fmt.Println("    cfs mount <region>:<fsid> <mountpoint>")
		os.Exit(1)
	}
	flag.Parse()
	if flag.NArg() == 0 {
		flag.Usage()
	}
	command := flag.Args()[0]
	switch command {
	default:
		fmt.Println("Invalid Command")
		flag.Usage()
	case "help":
		flag.Usage()
	case "version":
		fmt.Println("version:", cfsVersion)
		fmt.Println("commit:", commitVersion)
		fmt.Println("build date:", buildDate)
		fmt.Println("go version:", goVersion)
		os.Exit(0)
	case "configure":
		err := configure(configfile)
		if err != nil {
			fmt.Printf("Error writing config file: %v\n", err)
			os.Exit(1)
		}

	case "list":
		if !configured {
			fmt.Println("You must run \"cfs configure\" first.")
			os.Exit(1)
		}
		err := list(region, username, apikey)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	case "show":
		if !configured {
			fmt.Println("You must run \"cfs configure\" first.")
			os.Exit(1)
		}
		err := show(region, username, apikey)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	case "create":
		if !configured {
			fmt.Println("You must run \"cfs configure\" first.")
			os.Exit(1)
		}
		err := create(region, username, apikey)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	case "delete":
		if !configured {
			fmt.Println("You must run \"cfs configure\" first.")
			os.Exit(1)
		}
		err := del(region, username, apikey)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	case "update":
		if !configured {
			fmt.Println("You must run \"cfs configure\" first.")
			os.Exit(1)
		}
		err := update(region, username, apikey)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	case "grant":
		if !configured {
			fmt.Println("You must run \"cfs configure\" first.")
			os.Exit(1)
		}
		err := grant(region, username, apikey)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	case "revoke":
		if !configured {
			fmt.Println("You must run \"cfs configure\" first.")
			os.Exit(1)
		}
		err := grant(region, username, apikey)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	case "mount":
		err := mount()
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	}
}

func getArgs(args string) map[string]string {
	clargs := make(map[string]string)
	optList := strings.Split(args, ",")
	for _, item := range optList {
		if strings.Contains(item, "=") {
			value := strings.Split(item, "=")
			if value[0] == "" || value[1] == "" {
				log.Printf("Invalid option %s, %s no value\n\n", value[0], value[1])
				os.Exit(1)
			} else {
				clargs[value[0]] = value[1]
			}
		} else {
			clargs[item] = ""
		}
	}
	return clargs
}

// the mount command needs the path to fusermount
func fusermountPath() {
	// Grab the current path
	currentPath := os.Getenv("PATH")
	if len(currentPath) == 0 {
		// fusermount location for suse  /usr/bin
		// fusermount location on debian based distros /bin
		os.Setenv("PATH", "/usr/bin:/bin")
	}
}

// setupWS ...
func setupWS(svr string) *grpc.ClientConn {
	var opts []grpc.DialOption
	creds := credentials.NewTLS(&tls.Config{
		InsecureSkipVerify: true,
	})
	opts = append(opts, grpc.WithTransportCredentials(creds))
	conn, err := grpc.Dial(svr, opts...)
	if err != nil {
		log.Fatalf("failed to dial: %v", err)
	}
	return conn
}
<commit_msg>Change grant to revoke in the case statement<commit_after>package main

import (
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"os/user"
	"strings"
	"sync"

	"github.com/getcfs/fuse"
	pb "github.com/getcfs/megacfs/formic/proto"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
)

type server struct {
	fs *fs
	wg sync.WaitGroup
}

func newserver(fs *fs) *server {
	s := &server{
		fs: fs,
	}
	return s
}

func (s *server) serve() error {
	defer s.wg.Wait()

	for {
		req, err := s.fs.conn.ReadRequest()
		if err != nil {
			if err == io.EOF {
				break
			}
			return err
		}
		s.wg.Add(1)
		go func() {
			defer s.wg.Done()
			s.fs.handle(req)
		}()
	}
	return nil
}

func debuglog(msg interface{}) {
	fmt.Fprintf(os.Stderr, "%v\n", msg)
}

type rpc struct {
	conn *grpc.ClientConn
	api  pb.ApiClient
}

func newrpc(conn *grpc.ClientConn) *rpc {
	r := &rpc{
		conn: conn,
		api:  pb.NewApiClient(conn),
	}

	return r
}

// NullWriter ...
type NullWriter int

func (NullWriter) Write([]byte) (int, error) { return 0, nil }

var cfsVersion string
var buildDate string
var commitVersion string
var goVersion string

type authResponse struct {
	Access struct {
		Token struct {
			ID string `json:"id"`
		} `json:"token"`
	} `json:"access"`
}

func auth(username string, apikey string) string {
	body := fmt.Sprintf(`{"auth":{"RAX-KSKEY:apiKeyCredentials":{"username": "%s","apiKey": "%s"}}}`, username, apikey)
	rbody := strings.NewReader(body)
	req, err := http.NewRequest("POST", "https://identity.api.rackspacecloud.com/v2.0/tokens", rbody)
	if err != nil {
		fmt.Printf("%v", err)
		os.Exit(1)
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		fmt.Printf("%v", err)
		os.Exit(1)
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		fmt.Println(resp.Status)
		os.Exit(1)
	}

	// parse token from response
	var authResp authResponse
	r, _ := ioutil.ReadAll(resp.Body)
	json.Unmarshal(r, &authResp)
	token := authResp.Access.Token.ID

	return token
}

func mount() error {
	f := flag.NewFlagSet("mount", flag.ContinueOnError)
	f.Usage = func() {
		fmt.Println("Usage:")
		fmt.Println("    cfs mount [-o option,...] <region>:<fsid> <mountpoint>")
		fmt.Println("Options:")
		fmt.Println("    -o debug          enables debug output")
		fmt.Println("    -o ro             mount the filesystem read only")
		fmt.Println("    -o allow_other    allow access to other users")
		fmt.Println("Examples:")
		fmt.Println("    cfs mount iad:11111111-1111-1111-1111-111111111111 /mnt/test")
		fmt.Println("    cfs mount -o debug,ro iad:11111111-1111-1111-1111-111111111111 /mnt/test")
		os.Exit(1)
	}
	var options string
	f.StringVar(&options, "o", "", "")
	f.Parse(flag.Args()[1:])
	if f.NArg() != 2 {
		f.Usage()
	}
	regionFsid := f.Args()[0]
	parts := strings.Split(regionFsid, ":")
	if len(parts) != 2 {
		fmt.Println("Invalid filesystem:", regionFsid)
		f.Usage()
	}
	region := strings.ToLower(parts[0])
	fsid := parts[1]
	mountpoint := f.Args()[1]
	// Verify mountpoint exists
	_, err := os.Stat(mountpoint)
	if os.IsNotExist(err) {
		fmt.Printf("Mount point %s does not exist\n", mountpoint)
		os.Exit(1)
	}
	// Verify region
	addr, ok := regions[region]
	if !ok {
		fmt.Println("Invalid region:", region)
		os.Exit(1)
	}

	// Setup grpc
	var opts []grpc.DialOption
	creds := credentials.NewTLS(&tls.Config{
		InsecureSkipVerify: true,
	})
	opts = append(opts, grpc.WithTransportCredentials(creds))
	conn, err := grpc.Dial(addr, opts...)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	defer conn.Close()

	// handle fuse mount options
	mountOptions := []fuse.MountOption{
		fuse.FSName("cfs"),
		fuse.Subtype("cfs"),
		fuse.DefaultPermissions(),
		fuse.MaxReadahead(128 * 1024),
		fuse.AsyncRead(),
		//fuse.WritebackCache(),
		fuse.AutoInvalData(),
	}

	// parse mount options string
	clargs := getArgs(options)

	// handle debug mount option
	_, debug := clargs["debug"]
	if !debug {
		log.SetFlags(0)
		log.SetOutput(ioutil.Discard)
	}

	// handle allow_other mount option
	_, allowOther := clargs["allow_other"]
	if allowOther {
		mountOptions = append(mountOptions, fuse.AllowOther())
	}

	// handle ro mount option
	_, readOnly := clargs["ro"]
	if readOnly {
		mountOptions = append(mountOptions, fuse.ReadOnly())
	}

	// perform fuse mount
	fusermountPath()
	cfs, err := fuse.Mount(mountpoint, mountOptions...)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	defer cfs.Close()

	// setup rpc client
	rpc := newrpc(conn)
	fs := newfs(cfs, rpc, fsid)
	err = fs.InitFs()
	if err != nil {
		fmt.Println(err)
		exec.Command("fusermount", "-uz", mountpoint).Run()
		os.Exit(1)
	}
	srv := newserver(fs)

	if err := srv.serve(); err != nil {
		fmt.Println(err)
		exec.Command("fusermount", "-uz", mountpoint).Run()
		os.Exit(1)
	}

	<-cfs.Ready
	if err := cfs.MountError; err != nil {
		fmt.Println(err)
		exec.Command("fusermount", "-uz", mountpoint).Run()
		os.Exit(1)
	}
	return nil
}

func main() {

	// try to read config file
	configured := false
	config := map[string]string{}
	user, err := user.Current()
	if err != nil {
		fmt.Printf("Unable to identify curent user: %v\n", err)
		os.Exit(1)
	}
	var configfile string
	if user.HomeDir != "" {
		configfile = user.HomeDir + "/.cfs.json"
	} else {
		configfile = ".cfs.json"
	}
	f, err := ioutil.ReadFile(configfile)
	if err == nil {
		json.Unmarshal(f, &config)
		configured = true
	}
	region, _ := config["region"]
	username, _ := config["username"]
	apikey, _ := config["apikey"]

	flag.Usage = func() {
		fmt.Println("Usage:")
		fmt.Println("    cfs <command> [options] <args>")
		fmt.Println("Commands:")
		fmt.Println("    configure    create configuration file")
		fmt.Println("    list         list all filesystems")
		fmt.Println("    create       create a new filesytem")
		fmt.Println("    show         show filesystem details")
		fmt.Println("    update       update filesystem name")
		fmt.Println("    delete       delete an existing filesystem")
		fmt.Println("    grant        grant access to a filesystem")
		fmt.Println("    revoke       revoke access to a filesystem")
		fmt.Println("    mount        mount an existing filesystem")
		fmt.Println("    version      show client version")
		fmt.Println("    help         show usage for cfs")
		fmt.Println("Examples:")
		fmt.Println("    cfs configure")
		fmt.Println("    cfs create <name>")
		fmt.Println("    cfs grant <ip> <fsid>")
		fmt.Println("    cfs mount <region>:<fsid> <mountpoint>")
		os.Exit(1)
	}
	flag.Parse()
	if flag.NArg() == 0 {
		flag.Usage()
	}
	command := flag.Args()[0]
	switch command {
	default:
		fmt.Println("Invalid Command")
		flag.Usage()
	case "help":
		flag.Usage()
	case "version":
		fmt.Println("version:", cfsVersion)
		fmt.Println("commit:", commitVersion)
		fmt.Println("build date:", buildDate)
		fmt.Println("go version:", goVersion)
		os.Exit(0)
	case "configure":
		err := configure(configfile)
		if err != nil {
			fmt.Printf("Error writing config file: %v\n", err)
			os.Exit(1)
		}

	case "list":
		if !configured {
			fmt.Println("You must run \"cfs configure\" first.")
			os.Exit(1)
		}
		err := list(region, username, apikey)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	case "show":
		if !configured {
			fmt.Println("You must run \"cfs configure\" first.")
			os.Exit(1)
		}
		err := show(region, username, apikey)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	case "create":
		if !configured {
			fmt.Println("You must run \"cfs configure\" first.")
			os.Exit(1)
		}
		err := create(region, username, apikey)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	case "delete":
		if !configured {
			fmt.Println("You must run \"cfs configure\" first.")
			os.Exit(1)
		}
		err := del(region, username, apikey)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	case "update":
		if !configured {
			fmt.Println("You must run \"cfs configure\" first.")
			os.Exit(1)
		}
		err := update(region, username, apikey)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	case "grant":
		if !configured {
			fmt.Println("You must run \"cfs configure\" first.")
			os.Exit(1)
		}
		err := grant(region, username, apikey)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	case "revoke":
		if !configured {
			fmt.Println("You must run \"cfs configure\" first.")
			os.Exit(1)
		}
		err := revoke(region, username, apikey)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	case "mount":
		err := mount()
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	}
}

func getArgs(args string) map[string]string {
	clargs := make(map[string]string)
	optList := strings.Split(args, ",")
	for _, item := range optList {
		if strings.Contains(item, "=") {
			value := strings.Split(item, "=")
			if value[0] == "" || value[1] == "" {
				log.Printf("Invalid option %s, %s no value\n\n", value[0], value[1])
				os.Exit(1)
			} else {
				clargs[value[0]] = value[1]
			}
		} else {
			clargs[item] = ""
		}
	}
	return clargs
}

// the mount command needs the path to fusermount
func fusermountPath() {
	// Grab the current path
	currentPath := os.Getenv("PATH")
	if len(currentPath) == 0 {
		// fusermount location for suse  /usr/bin
		// fusermount location on debian based distros /bin
		os.Setenv("PATH", "/usr/bin:/bin")
	}
}

// setupWS ...
func setupWS(svr string) *grpc.ClientConn {
	var opts []grpc.DialOption
	creds := credentials.NewTLS(&tls.Config{
		InsecureSkipVerify: true,
	})
	opts = append(opts, grpc.WithTransportCredentials(creds))
	conn, err := grpc.Dial(svr, opts...)
	if err != nil {
		log.Fatalf("failed to dial: %v", err)
	}
	return conn
}
<|endoftext|>


<commit_before>// Package boilingcore has types and methods useful for generating code that
// acts as a fully dynamic ORM might.
package boilingcore

import (
	"encoding/json"
	"fmt"
	"go/build"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"text/template"

	"github.com/pkg/errors"
	"github.com/vattle/sqlboiler/bdb"
	"github.com/vattle/sqlboiler/bdb/drivers"
	"github.com/vattle/sqlboiler/queries"
	"github.com/vattle/sqlboiler/strmangle"
)

const (
	templatesDirectory          = "templates"
	templatesSingletonDirectory = "templates/singleton"

	templatesTestDirectory          = "templates_test"
	templatesSingletonTestDirectory = "templates_test/singleton"

	templatesTestMainDirectory = "templates_test/main_test"
)

// State holds the global data needed by most pieces to run
type State struct {
	Config *Config

	Driver  bdb.Interface
	Tables  []bdb.Table
	Dialect queries.Dialect

	Templates              *templateList
	TestTemplates          *templateList
	SingletonTemplates     *templateList
	SingletonTestTemplates *templateList

	TestMainTemplate *template.Template
}

// New creates a new state based off of the config
func New(config *Config) (*State, error) {
	s := &State{
		Config: config,
	}

	err := s.initDriver(config.DriverName)
	if err != nil {
		return nil, err
	}

	// Connect to the driver database
	if err = s.Driver.Open(); err != nil {
		return nil, errors.Wrap(err, "unable to connect to the database")
	}

	err = s.initTables(config.Schema, config.WhitelistTables, config.BlacklistTables)
	if err != nil {
		return nil, errors.Wrap(err, "unable to initialize tables")
	}

	if s.Config.Debug {
		b, err := json.Marshal(s.Tables)
		if err != nil {
			return nil, errors.Wrap(err, "unable to json marshal tables")
		}
		fmt.Printf("%s\n", b)
	}

	err = s.initOutFolder()
	if err != nil {
		return nil, errors.Wrap(err, "unable to initialize the output folder")
	}

	err = s.initTemplates()
	if err != nil {
		return nil, errors.Wrap(err, "unable to initialize templates")
	}

	err = s.initTags(config.Tags)
	if err != nil {
		return nil, errors.Wrap(err, "unable to initialize struct tags")
	}

	return s, nil
}

// Run executes the sqlboiler templates and outputs them to files based on the
// state given.
func (s *State) Run(includeTests bool) error {
	singletonData := &templateData{
		Tables:           s.Tables,
		Schema:           s.Config.Schema,
		DriverName:       s.Config.DriverName,
		UseLastInsertID:  s.Driver.UseLastInsertID(),
		PkgName:          s.Config.PkgName,
		NoHooks:          s.Config.NoHooks,
		NoAutoTimestamps: s.Config.NoAutoTimestamps,
		Dialect:          s.Dialect,
		LQ:               strmangle.QuoteCharacter(s.Dialect.LQ),
		RQ:               strmangle.QuoteCharacter(s.Dialect.RQ),

		StringFuncs: templateStringMappers,
	}

	if err := generateSingletonOutput(s, singletonData); err != nil {
		return errors.Wrap(err, "singleton template output")
	}

	if !s.Config.NoTests && includeTests {
		if err := generateTestMainOutput(s, singletonData); err != nil {
			return errors.Wrap(err, "unable to generate TestMain output")
		}

		if err := generateSingletonTestOutput(s, singletonData); err != nil {
			return errors.Wrap(err, "unable to generate singleton test template output")
		}
	}

	for _, table := range s.Tables {
		if table.IsJoinTable {
			continue
		}

		data := &templateData{
			Tables:           s.Tables,
			Table:            table,
			Schema:           s.Config.Schema,
			DriverName:       s.Config.DriverName,
			UseLastInsertID:  s.Driver.UseLastInsertID(),
			PkgName:          s.Config.PkgName,
			NoHooks:          s.Config.NoHooks,
			NoAutoTimestamps: s.Config.NoAutoTimestamps,
			Tags:             s.Config.Tags,
			Dialect:          s.Dialect,
			LQ:               strmangle.QuoteCharacter(s.Dialect.LQ),
			RQ:               strmangle.QuoteCharacter(s.Dialect.RQ),

			StringFuncs: templateStringMappers,
		}

		// Generate the regular templates
		if err := generateOutput(s, data); err != nil {
			return errors.Wrap(err, "unable to generate output")
		}

		// Generate the test templates
		if !s.Config.NoTests && includeTests {
			if err := generateTestOutput(s, data); err != nil {
				return errors.Wrap(err, "unable to generate test output")
			}
		}
	}

	return nil
}

// Cleanup closes any resources that must be closed
func (s *State) Cleanup() error {
	s.Driver.Close()
	return nil
}

// initTemplates loads all template folders into the state object.
func (s *State) initTemplates() error {
	var err error

	basePath, err := getBasePath(s.Config.BaseDir)
	if err != nil {
		return err
	}

	s.Templates, err = loadTemplates(filepath.Join(basePath, templatesDirectory))
	if err != nil {
		return err
	}

	s.SingletonTemplates, err = loadTemplates(filepath.Join(basePath, templatesSingletonDirectory))
	if err != nil {
		return err
	}

	if !s.Config.NoTests {
		s.TestTemplates, err = loadTemplates(filepath.Join(basePath, templatesTestDirectory))
		if err != nil {
			return err
		}

		s.SingletonTestTemplates, err = loadTemplates(filepath.Join(basePath, templatesSingletonTestDirectory))
		if err != nil {
			return err
		}

		testMain := s.Config.DriverName + "_main.tpl"
		s.TestMainTemplate, err = loadTemplate(nil, testMain, filepath.Join(basePath, templatesTestMainDirectory, testMain))
		if err != nil {
			return err
		}
	}

	return s.processReplacements()
}

// processReplacements loads any replacement templates
func (s *State) processReplacements() error {
	for _, replace := range s.Config.Replacements {
		fmt.Println(replace)
		splits := strings.Split(replace, ":")
		if len(splits) != 2 {
			return errors.Errorf("replace parameters must have 2 arguments, given: %s", replace)
		}

		toReplace, replaceWith := splits[0], splits[1]

		var err error
		switch filepath.Dir(toReplace) {
		case templatesDirectory:
			_, err = loadTemplate(s.Templates.Template, toReplace, replaceWith)
		case templatesSingletonDirectory:
			_, err = loadTemplate(s.SingletonTemplates.Template, toReplace, replaceWith)
		case templatesTestDirectory:
			_, err = loadTemplate(s.TestTemplates.Template, toReplace, replaceWith)
		case templatesSingletonTestDirectory:
			_, err = loadTemplate(s.SingletonTestTemplates.Template, toReplace, replaceWith)
		case templatesTestMainDirectory:
			s.TestMainTemplate, err = loadTemplate(nil, toReplace, replaceWith)
		default:
			return errors.Errorf("replace file's directory not part of any known folder: %s", toReplace)
		}

		if err != nil {
			return err
		}
	}

	return nil
}

var basePackage = "github.com/vattle/sqlboiler"

func getBasePath(baseDirConfig string) (string, error) {
	if len(baseDirConfig) > 0 {
		return baseDirConfig, nil
	}

	p, _ := build.Default.Import(basePackage, "", build.FindOnly)
	if p != nil && len(p.Dir) > 0 {
		return p.Dir, nil
	}

	return os.Getwd()
}

// initDriver attempts to set the state Interface based off the passed in
// driver flag value. If an invalid flag string is provided an error is returned.
func (s *State) initDriver(driverName string) error {
	// Create a driver based off driver flag
	switch driverName {
	case "postgres":
		s.Driver = drivers.NewPostgresDriver(
			s.Config.Postgres.User,
			s.Config.Postgres.Pass,
			s.Config.Postgres.DBName,
			s.Config.Postgres.Host,
			s.Config.Postgres.Port,
			s.Config.Postgres.SSLMode,
		)
	case "mysql":
		s.Driver = drivers.NewMySQLDriver(
			s.Config.MySQL.User,
			s.Config.MySQL.Pass,
			s.Config.MySQL.DBName,
			s.Config.MySQL.Host,
			s.Config.MySQL.Port,
			s.Config.MySQL.SSLMode,
		)
	case "mock":
		s.Driver = &drivers.MockDriver{}
	}

	if s.Driver == nil {
		return errors.New("An invalid driver name was provided")
	}

	s.Dialect.LQ = s.Driver.LeftQuote()
	s.Dialect.RQ = s.Driver.RightQuote()
	s.Dialect.IndexPlaceholders = s.Driver.IndexPlaceholders()

	return nil
}

// initTables retrieves all "public" schema table names from the database.
func (s *State) initTables(schema string, whitelist, blacklist []string) error {
	var err error
	s.Tables, err = bdb.Tables(s.Driver, schema, whitelist, blacklist)
	if err != nil {
		return errors.Wrap(err, "unable to fetch table data")
	}

	if len(s.Tables) == 0 {
		return errors.New("no tables found in database")
	}

	if err := checkPKeys(s.Tables); err != nil {
		return err
	}

	return nil
}

// Tags must be in a format like: json, xml, etc.
var rgxValidTag = regexp.MustCompile(`[a-zA-Z_\.]+`)

// initTags removes duplicate tags and validates the format
// of all user tags are simple strings without quotes: [a-zA-Z_\.]+
func (s *State) initTags(tags []string) error {
	s.Config.Tags = removeDuplicates(s.Config.Tags)
	for _, v := range s.Config.Tags {
		if !rgxValidTag.MatchString(v) {
			return errors.New("Invalid tag format %q supplied, only specify name, eg: xml")
		}
	}

	return nil
}

// initOutFolder creates the folder that will hold the generated output.
func (s *State) initOutFolder() error {
	if s.Config.Wipe {
		if err := os.RemoveAll(s.Config.OutFolder); err != nil {
			return err
		}
	}

	return os.MkdirAll(s.Config.OutFolder, os.ModePerm)
}

// checkPKeys ensures every table has a primary key column
func checkPKeys(tables []bdb.Table) error {
	var missingPkey []string
	for _, t := range tables {
		if t.PKey == nil {
			missingPkey = append(missingPkey, t.Name)
		}
	}

	if len(missingPkey) != 0 {
		return errors.Errorf("primary key missing in tables (%s)", strings.Join(missingPkey, ", "))
	}

	return nil
}
<commit_msg>Remove println that should not be present<commit_after>// Package boilingcore has types and methods useful for generating code that
// acts as a fully dynamic ORM might.
package boilingcore

import (
	"encoding/json"
	"fmt"
	"go/build"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"text/template"

	"github.com/pkg/errors"
	"github.com/vattle/sqlboiler/bdb"
	"github.com/vattle/sqlboiler/bdb/drivers"
	"github.com/vattle/sqlboiler/queries"
	"github.com/vattle/sqlboiler/strmangle"
)

const (
	templatesDirectory          = "templates"
	templatesSingletonDirectory = "templates/singleton"

	templatesTestDirectory          = "templates_test"
	templatesSingletonTestDirectory = "templates_test/singleton"

	templatesTestMainDirectory = "templates_test/main_test"
)

// State holds the global data needed by most pieces to run
type State struct {
	Config *Config

	Driver  bdb.Interface
	Tables  []bdb.Table
	Dialect queries.Dialect

	Templates              *templateList
	TestTemplates          *templateList
	SingletonTemplates     *templateList
	SingletonTestTemplates *templateList

	TestMainTemplate *template.Template
}

// New creates a new state based off of the config
func New(config *Config) (*State, error) {
	s := &State{
		Config: config,
	}

	err := s.initDriver(config.DriverName)
	if err != nil {
		return nil, err
	}

	// Connect to the driver database
	if err = s.Driver.Open(); err != nil {
		return nil, errors.Wrap(err, "unable to connect to the database")
	}

	err = s.initTables(config.Schema, config.WhitelistTables, config.BlacklistTables)
	if err != nil {
		return nil, errors.Wrap(err, "unable to initialize tables")
	}

	if s.Config.Debug {
		b, err := json.Marshal(s.Tables)
		if err != nil {
			return nil, errors.Wrap(err, "unable to json marshal tables")
		}
		fmt.Printf("%s\n", b)
	}

	err = s.initOutFolder()
	if err != nil {
		return nil, errors.Wrap(err, "unable to initialize the output folder")
	}

	err = s.initTemplates()
	if err != nil {
		return nil, errors.Wrap(err, "unable to initialize templates")
	}

	err = s.initTags(config.Tags)
	if err != nil {
		return nil, errors.Wrap(err, "unable to initialize struct tags")
	}

	return s, nil
}

// Run executes the sqlboiler templates and outputs them to files based on the
// state given.
func (s *State) Run(includeTests bool) error {
	singletonData := &templateData{
		Tables:           s.Tables,
		Schema:           s.Config.Schema,
		DriverName:       s.Config.DriverName,
		UseLastInsertID:  s.Driver.UseLastInsertID(),
		PkgName:          s.Config.PkgName,
		NoHooks:          s.Config.NoHooks,
		NoAutoTimestamps: s.Config.NoAutoTimestamps,
		Dialect:          s.Dialect,
		LQ:               strmangle.QuoteCharacter(s.Dialect.LQ),
		RQ:               strmangle.QuoteCharacter(s.Dialect.RQ),

		StringFuncs: templateStringMappers,
	}

	if err := generateSingletonOutput(s, singletonData); err != nil {
		return errors.Wrap(err, "singleton template output")
	}

	if !s.Config.NoTests && includeTests {
		if err := generateTestMainOutput(s, singletonData); err != nil {
			return errors.Wrap(err, "unable to generate TestMain output")
		}

		if err := generateSingletonTestOutput(s, singletonData); err != nil {
			return errors.Wrap(err, "unable to generate singleton test template output")
		}
	}

	for _, table := range s.Tables {
		if table.IsJoinTable {
			continue
		}

		data := &templateData{
			Tables:           s.Tables,
			Table:            table,
			Schema:           s.Config.Schema,
			DriverName:       s.Config.DriverName,
			UseLastInsertID:  s.Driver.UseLastInsertID(),
			PkgName:          s.Config.PkgName,
			NoHooks:          s.Config.NoHooks,
			NoAutoTimestamps: s.Config.NoAutoTimestamps,
			Tags:             s.Config.Tags,
			Dialect:          s.Dialect,
			LQ:               strmangle.QuoteCharacter(s.Dialect.LQ),
			RQ:               strmangle.QuoteCharacter(s.Dialect.RQ),

			StringFuncs: templateStringMappers,
		}

		// Generate the regular templates
		if err := generateOutput(s, data); err != nil {
			return errors.Wrap(err, "unable to generate output")
		}

		// Generate the test templates
		if !s.Config.NoTests && includeTests {
			if err := generateTestOutput(s, data); err != nil {
				return errors.Wrap(err, "unable to generate test output")
			}
		}
	}

	return nil
}

// Cleanup closes any resources that must be closed
func (s *State) Cleanup() error {
	s.Driver.Close()
	return nil
}

// initTemplates loads all template folders into the state object.
func (s *State) initTemplates() error {
	var err error

	basePath, err := getBasePath(s.Config.BaseDir)
	if err != nil {
		return err
	}

	s.Templates, err = loadTemplates(filepath.Join(basePath, templatesDirectory))
	if err != nil {
		return err
	}

	s.SingletonTemplates, err = loadTemplates(filepath.Join(basePath, templatesSingletonDirectory))
	if err != nil {
		return err
	}

	if !s.Config.NoTests {
		s.TestTemplates, err = loadTemplates(filepath.Join(basePath, templatesTestDirectory))
		if err != nil {
			return err
		}

		s.SingletonTestTemplates, err = loadTemplates(filepath.Join(basePath, templatesSingletonTestDirectory))
		if err != nil {
			return err
		}

		testMain := s.Config.DriverName + "_main.tpl"
		s.TestMainTemplate, err = loadTemplate(nil, testMain, filepath.Join(basePath, templatesTestMainDirectory, testMain))
		if err != nil {
			return err
		}
	}

	return s.processReplacements()
}

// processReplacements loads any replacement templates
func (s *State) processReplacements() error {
	for _, replace := range s.Config.Replacements {
		splits := strings.Split(replace, ":")
		if len(splits) != 2 {
			return errors.Errorf("replace parameters must have 2 arguments, given: %s", replace)
		}

		toReplace, replaceWith := splits[0], splits[1]

		var err error
		switch filepath.Dir(toReplace) {
		case templatesDirectory:
			_, err = loadTemplate(s.Templates.Template, toReplace, replaceWith)
		case templatesSingletonDirectory:
			_, err = loadTemplate(s.SingletonTemplates.Template, toReplace, replaceWith)
		case templatesTestDirectory:
			_, err = loadTemplate(s.TestTemplates.Template, toReplace, replaceWith)
		case templatesSingletonTestDirectory:
			_, err = loadTemplate(s.SingletonTestTemplates.Template, toReplace, replaceWith)
		case templatesTestMainDirectory:
			s.TestMainTemplate, err = loadTemplate(nil, toReplace, replaceWith)
		default:
			return errors.Errorf("replace file's directory not part of any known folder: %s", toReplace)
		}

		if err != nil {
			return err
		}
	}

	return nil
}

var basePackage = "github.com/vattle/sqlboiler"

func getBasePath(baseDirConfig string) (string, error) {
	if len(baseDirConfig) > 0 {
		return baseDirConfig, nil
	}

	p, _ := build.Default.Import(basePackage, "", build.FindOnly)
	if p != nil && len(p.Dir) > 0 {
		return p.Dir, nil
	}

	return os.Getwd()
}

// initDriver attempts to set the state Interface based off the passed in
// driver flag value. If an invalid flag string is provided an error is returned.
func (s *State) initDriver(driverName string) error {
	// Create a driver based off driver flag
	switch driverName {
	case "postgres":
		s.Driver = drivers.NewPostgresDriver(
			s.Config.Postgres.User,
			s.Config.Postgres.Pass,
			s.Config.Postgres.DBName,
			s.Config.Postgres.Host,
			s.Config.Postgres.Port,
			s.Config.Postgres.SSLMode,
		)
	case "mysql":
		s.Driver = drivers.NewMySQLDriver(
			s.Config.MySQL.User,
			s.Config.MySQL.Pass,
			s.Config.MySQL.DBName,
			s.Config.MySQL.Host,
			s.Config.MySQL.Port,
			s.Config.MySQL.SSLMode,
		)
	case "mock":
		s.Driver = &drivers.MockDriver{}
	}

	if s.Driver == nil {
		return errors.New("An invalid driver name was provided")
	}

	s.Dialect.LQ = s.Driver.LeftQuote()
	s.Dialect.RQ = s.Driver.RightQuote()
	s.Dialect.IndexPlaceholders = s.Driver.IndexPlaceholders()

	return nil
}

// initTables retrieves all "public" schema table names from the database.
func (s *State) initTables(schema string, whitelist, blacklist []string) error {
	var err error
	s.Tables, err = bdb.Tables(s.Driver, schema, whitelist, blacklist)
	if err != nil {
		return errors.Wrap(err, "unable to fetch table data")
	}

	if len(s.Tables) == 0 {
		return errors.New("no tables found in database")
	}

	if err := checkPKeys(s.Tables); err != nil {
		return err
	}

	return nil
}

// Tags must be in a format like: json, xml, etc.
var rgxValidTag = regexp.MustCompile(`[a-zA-Z_\.]+`)

// initTags removes duplicate tags and validates the format
// of all user tags are simple strings without quotes: [a-zA-Z_\.]+
func (s *State) initTags(tags []string) error {
	s.Config.Tags = removeDuplicates(s.Config.Tags)
	for _, v := range s.Config.Tags {
		if !rgxValidTag.MatchString(v) {
			return errors.New("Invalid tag format %q supplied, only specify name, eg: xml")
		}
	}

	return nil
}

// initOutFolder creates the folder that will hold the generated output.
func (s *State) initOutFolder() error {
	if s.Config.Wipe {
		if err := os.RemoveAll(s.Config.OutFolder); err != nil {
			return err
		}
	}

	return os.MkdirAll(s.Config.OutFolder, os.ModePerm)
}

// checkPKeys ensures every table has a primary key column
func checkPKeys(tables []bdb.Table) error {
	var missingPkey []string
	for _, t := range tables {
		if t.PKey == nil {
			missingPkey = append(missingPkey, t.Name)
		}
	}

	if len(missingPkey) != 0 {
		return errors.Errorf("primary key missing in tables (%s)", strings.Join(missingPkey, ", "))
	}

	return nil
}
<|endoftext|>


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

package io

import (
	"os";
	"syscall";
)

export var ErrEOF = os.NewError("EOF")

export type Read interface {
	Read(p *[]byte) (n int, err *os.Error);
}

export type Write interface {
	Write(p *[]byte) (n int, err *os.Error);
}

export type ReadWrite interface {
	Read(p *[]byte) (n int, err *os.Error);
	Write(p *[]byte) (n int, err *os.Error);
}

export type ReadWriteClose interface {
	Read(p *[]byte) (n int, err *os.Error);
	Write(p *[]byte) (n int, err *os.Error);
	Close() *os.Error;
}

export func WriteString(w Write, s string) (n int, err *os.Error) {
	b := new([]byte, len(s)+1);
	if !syscall.StringToBytes(b, s) {
		return -1, os.EINVAL
	}
	// BUG return w.Write(b[0:len(s)])
	r, e := w.Write(b[0:len(s)]);
	return r, e
}

// Read until buffer is full,  EOF, or error
export func Readn(fd Read, buf *[]byte) (n int, err *os.Error) {
	n = 0;
	for n < len(buf) {
		nn, e := fd.Read(buf[n:len(buf)]);
		if nn > 0 {
			n += nn
		}
		if e != nil {
			return n, e
		}
		if nn <= 0 {
			return n, ErrEOF	// no error but insufficient data
		}
	}
	return n, nil
}

// Convert something that implements Read into something
// whose Reads are always Readn
type FullRead struct {
	fd	Read;
}

func (fd *FullRead) Read(p *[]byte) (n int, err *os.Error) {
	n, err = Readn(fd.fd, p);
	return n, err
}

export func MakeFullReader(fd Read) Read {
	if fr, ok := fd.(*FullRead); ok {
		// already a FullRead
		return fd
	}
	return &FullRead{fd}
}
<commit_msg>Copyn<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package io

import (
	"os";
	"syscall";
)

export var ErrEOF = os.NewError("EOF")

export type Read interface {
	Read(p *[]byte) (n int, err *os.Error);
}

export type Write interface {
	Write(p *[]byte) (n int, err *os.Error);
}

export type ReadWrite interface {
	Read(p *[]byte) (n int, err *os.Error);
	Write(p *[]byte) (n int, err *os.Error);
}

export type ReadWriteClose interface {
	Read(p *[]byte) (n int, err *os.Error);
	Write(p *[]byte) (n int, err *os.Error);
	Close() *os.Error;
}

export func WriteString(w Write, s string) (n int, err *os.Error) {
	b := new([]byte, len(s)+1);
	if !syscall.StringToBytes(b, s) {
		return -1, os.EINVAL
	}
	// BUG return w.Write(b[0:len(s)])
	r, e := w.Write(b[0:len(s)]);
	return r, e
}

// Read until buffer is full,  EOF, or error
export func Readn(fd Read, buf *[]byte) (n int, err *os.Error) {
	n = 0;
	for n < len(buf) {
		nn, e := fd.Read(buf[n:len(buf)]);
		if nn > 0 {
			n += nn
		}
		if e != nil {
			return n, e
		}
		if nn <= 0 {
			return n, ErrEOF	// no error but insufficient data
		}
	}
	return n, nil
}

// Convert something that implements Read into something
// whose Reads are always Readn
type FullRead struct {
	fd	Read;
}

func (fd *FullRead) Read(p *[]byte) (n int, err *os.Error) {
	n, err = Readn(fd.fd, p);
	return n, err
}

export func MakeFullReader(fd Read) Read {
	if fr, ok := fd.(*FullRead); ok {
		// already a FullRead
		return fd
	}
	return &FullRead{fd}
}

// Copies n bytes (or until EOF is reached) from src to dst.
// Returns the number of bytes copied and the error, if any.
export func Copyn(src Read, dst Write, n int) (c int, err *os.Error) {
	buf := new([]byte, 32*1024);  // BUG 6g crashes on non-pointer array slices
	c = 0;
	for c < n {
		l := n - c;
		if l > len(buf) {
			l = len(buf)
		}
		nr, er := src.Read(buf[0 : l]);
		if nr > 0 {
			nw, ew := dst.Write(buf[0 : nr]);
			if nw != nr || ew != nil {
				c += nw;
				if ew == nil {
					ew = os.EIO
				}
				err = ew;
				break;
			}
			c += nr;
		}
		if er != nil {
			err = er;
			break;
		}
		if nr == 0 {
			break;
		}
	}
	return c, err
}
<|endoftext|>