<commit_before>// Copyright 2018 Google LLC
//
// 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 httpreplay_test

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"
	"time"

	"cloud.google.com/go/httpreplay"
	"cloud.google.com/go/internal/testutil"
	"cloud.google.com/go/storage"
	"google.golang.org/api/option"
)

func TestIntegration_RecordAndReplay(t *testing.T) {
	httpreplay.DebugHeaders()
	if testing.Short() {
		t.Skip("Integration tests skipped in short mode")
	}
	replayFilename := tempFilename(t, "RecordAndReplay*.replay")
	defer os.Remove(replayFilename)
	projectID := testutil.ProjID()
	if projectID == "" {
		t.Skip("Need project ID. See CONTRIBUTING.md for details.")
	}
	ctx := context.Background()

	// Record.
	initial := time.Now()
	ibytes, err := json.Marshal(initial)
	if err != nil {
		t.Fatal(err)
	}
	rec, err := httpreplay.NewRecorder(replayFilename, ibytes)
	if err != nil {
		t.Fatal(err)
	}
	hc, err := rec.Client(ctx, option.WithTokenSource(
		testutil.TokenSource(ctx, storage.ScopeFullControl)))
	if err != nil {
		t.Fatal(err)
	}
	wanta, wantc := run(t, hc)
	testReadCRC(t, hc, "recording")
	if err := rec.Close(); err != nil {
		t.Fatalf("rec.Close: %v", err)
	}

	// Replay.
	rep, err := httpreplay.NewReplayer(replayFilename)
	if err != nil {
		t.Fatal(err)
	}
	defer rep.Close()
	hc, err = rep.Client(ctx)
	if err != nil {
		t.Fatal(err)
	}
	gota, gotc := run(t, hc)
	testReadCRC(t, hc, "replaying")

	if diff := testutil.Diff(gota, wanta); diff != "" {
		t.Error(diff)
	}
	if !bytes.Equal(gotc, wantc) {
		t.Errorf("got %q, want %q", gotc, wantc)
	}
	var gotInitial time.Time
	if err := json.Unmarshal(rep.Initial(), &gotInitial); err != nil {
		t.Fatal(err)
	}
	if !gotInitial.Equal(initial) {
		t.Errorf("initial: got %v, want %v", gotInitial, initial)
	}
}

// TODO(jba): test errors

func run(t *testing.T, hc *http.Client) (*storage.BucketAttrs, []byte) {
	ctx := context.Background()
	client, err := storage.NewClient(ctx, option.WithHTTPClient(hc))
	if err != nil {
		t.Fatal(err)
	}
	defer client.Close()
	b := client.Bucket(testutil.ProjID())
	attrs, err := b.Attrs(ctx)
	if err != nil {
		t.Fatal(err)
	}
	obj := b.Object("replay-test")
	w := obj.NewWriter(ctx)
	data := []byte{150, 151, 152}
	if _, err := w.Write(data); err != nil {
		t.Fatal(err)
	}
	if err := w.Close(); err != nil {
		t.Fatal(err)
	}

	r, err := obj.NewReader(ctx)
	if err != nil {
		t.Fatal(err)
	}
	defer r.Close()
	contents, err := ioutil.ReadAll(r)
	if err != nil {
		t.Fatal(err)
	}

	return attrs, contents
}

func testReadCRC(t *testing.T, hc *http.Client, mode string) {
	const (
		// This is an uncompressed file.
		// See https://cloud.google.com/storage/docs/public-datasets/landsat
		uncompressedBucket = "gcp-public-data-landsat"
		uncompressedObject = "LC08/PRE/044/034/LC80440342016259LGN00/LC80440342016259LGN00_MTL.txt"

		gzippedBucket = "storage-library-test-bucket"
		gzippedObject = "gzipped-text.txt"
	)
	ctx := context.Background()
	client, err := storage.NewClient(ctx, option.WithHTTPClient(hc))
	if err != nil {
		t.Fatalf("%s: %v", mode, err)
	}
	defer client.Close()

	uncompressedObj := client.Bucket(uncompressedBucket).Object(uncompressedObject)
	gzippedObj := client.Bucket(gzippedBucket).Object(gzippedObject)

	for _, test := range []struct {
		desc           string
		obj            *storage.ObjectHandle
		offset, length int64
		readCompressed bool // don't decompress a gzipped file

		wantErr bool
		wantLen int // length of contents
	}{
		{
			desc:           "uncompressed, entire file",
			obj:            uncompressedObj,
			offset:         0,
			length:         -1,
			readCompressed: false,
			wantLen:        7903,
		},
		{
			desc:           "uncompressed, entire file, don't decompress",
			obj:            uncompressedObj,
			offset:         0,
			length:         -1,
			readCompressed: true,
			wantLen:        7903,
		},
		{
			desc:           "uncompressed, suffix",
			obj:            uncompressedObj,
			offset:         3,
			length:         -1,
			readCompressed: false,
			wantLen:        7900,
		},
		{
			desc:           "uncompressed, prefix",
			obj:            uncompressedObj,
			offset:         0,
			length:         18,
			readCompressed: false,
			wantLen:        18,
		},
		{
			// When a gzipped file is unzipped by GCS, we can't verify the checksum
			// because it was computed against the zipped contents. There is no
			// header that indicates that a gzipped file is being served unzipped.
			// But our CRC check only happens if there is a Content-Length header,
			// and that header is absent for this read.
			desc:           "compressed, entire file, server unzips",
			obj:            gzippedObj,
			offset:         0,
			length:         -1,
			readCompressed: false,
			wantLen:        11,
		},
		{
			// When we read a gzipped file uncompressed, it's like reading a regular file:
			// the served content and the CRC match.
			desc:           "compressed, entire file, read compressed",
			obj:            gzippedObj,
			offset:         0,
			length:         -1,
			readCompressed: true,
			wantLen:        31,
		},
		{
			desc:           "compressed, partial, read compressed",
			obj:            gzippedObj,
			offset:         1,
			length:         8,
			readCompressed: true,
			wantLen:        8,
		},
		{
			desc:    "uncompressed, HEAD",
			obj:     uncompressedObj,
			offset:  0,
			length:  0,
			wantLen: 0,
		},
		{
			desc:    "compressed, HEAD",
			obj:     gzippedObj,
			offset:  0,
			length:  0,
			wantLen: 0,
		},
	} {
		obj := test.obj.ReadCompressed(test.readCompressed)
		r, err := obj.NewRangeReader(ctx, test.offset, test.length)
		if err != nil {
			if test.wantErr {
				continue
			}
			t.Errorf("%s: %s: %v", mode, test.desc, err)
			continue
		}
		data, err := ioutil.ReadAll(r)
		_ = r.Close()
		if err != nil {
			t.Errorf("%s: %s: %v", mode, test.desc, err)
			continue
		}
		if got, want := len(data), test.wantLen; got != want {
			t.Errorf("%s: %s: len: got %d, want %d", mode, test.desc, got, want)
		}
	}
}

func TestRemoveAndClear(t *testing.T) {
	if testing.Short() {
		t.Skip("Integration tests skipped in short mode")
	}
	// Disable logging for this test, since it generates a lot.
	log.SetOutput(ioutil.Discard)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		fmt.Fprintln(w, "LGTM")
	}))
	defer srv.Close()

	replayFilename := tempFilename(t, "TestRemoveAndClear*.replay")
	defer os.Remove(replayFilename)

	ctx := context.Background()
	// Record
	rec, err := httpreplay.NewRecorder(replayFilename, nil)
	if err != nil {
		t.Fatal(err)
	}
	rec.ClearHeaders("Clear")
	rec.RemoveRequestHeaders("Rem*")
	rec.ClearQueryParams("c")
	rec.RemoveQueryParams("r")
	hc, err := rec.Client(ctx)
	if err != nil {
		t.Fatal(err)
	}
	query := "k=1&r=2&c=3"
	req, err := http.NewRequest("GET", srv.URL+"?"+query, nil)
	if err != nil {
		t.Fatal(err)
	}
	headers := map[string]string{"Keep": "ok", "Clear": "secret", "Remove": "bye"}
	for k, v := range headers {
		req.Header.Set(k, v)
	}
	if _, err := hc.Do(req); err != nil {
		t.Fatal(err)
	}
	if err := rec.Close(); err != nil {
		t.Fatal(err)
	}

	// Replay
	// For both headers and query param:
	// - k or Keep must be present and identical
	// - c or Clear must be present, but can be different
	// - r or Remove can be anything
	for _, test := range []struct {
		query       string
		headers     map[string]string
		wantSuccess bool
	}{
		{query, headers, true}, // same query string and headers
		{query,
			map[string]string{"Keep": "oops", "Clear": "secret", "Remove": "bye"},
			false, // different Keep
		},
		{query, map[string]string{}, false},                               // missing Keep and Clear
		{query, map[string]string{"Keep": "ok"}, false},                   // missing Clear
		{query, map[string]string{"Keep": "ok", "Clear": "secret"}, true}, // missing Remove is OK
		{
			query,
			map[string]string{"Keep": "ok", "Clear": "secret", "Remove": "whatev"},
			true,
		}, // different Remove is OK
		{query, map[string]string{"Keep": "ok", "Clear": "diff"}, true}, // different Clear is OK
		{"", headers, false},            // no query string
		{"k=x&r=2&c=3", headers, false}, // different k
		{"r=2", headers, false},         // missing k and c
		{"k=1&r=2", headers, false},     // missing c
		{"k=1&c=3", headers, true},      // missing r is OK
		{"k=1&r=x&c=3", headers, true},  // different r is OK,
		{"k=1&r=2&c=x", headers, true},  // different clear is OK
	} {
		rep, err := httpreplay.NewReplayer(replayFilename)
		if err != nil {
			t.Fatal(err)
		}
		hc, err = rep.Client(ctx)
		if err != nil {
			t.Fatal(err)
		}
		url := srv.URL
		if test.query != "" {
			url += "?" + test.query
		}
		req, err = http.NewRequest("GET", url, nil)
		if err != nil {
			t.Fatal(err)
		}
		for k, v := range test.headers {
			req.Header.Set(k, v)
		}
		resp, err := hc.Do(req)
		if err != nil {
			t.Fatal(err)
		}
		rep.Close()
		if (resp.StatusCode == 200) != test.wantSuccess {
			t.Errorf("%q, %v: got %d, wanted success=%t",
				test.query, test.headers, resp.StatusCode, test.wantSuccess)
		}
	}
}

func tempFilename(t *testing.T, pattern string) string {
	f, err := ioutil.TempFile("", pattern)
	if err != nil {
		t.Fatal(err)
	}
	filename := f.Name()
	if err := f.Close(); err != nil {
		t.Fatal(err)
	}
	return filename
}
<commit_msg>httpreplay: fix test for kokoro<commit_after>// Copyright 2018 Google LLC
//
// 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 httpreplay_test

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"
	"time"

	"cloud.google.com/go/httpreplay"
	"cloud.google.com/go/internal/testutil"
	"cloud.google.com/go/storage"
	"google.golang.org/api/option"
)

func TestIntegration_RecordAndReplay(t *testing.T) {
	httpreplay.DebugHeaders()
	if testing.Short() {
		t.Skip("Integration tests skipped in short mode")
	}
	replayFilename := tempFilename(t, "RecordAndReplay*.replay")
	defer os.Remove(replayFilename)
	projectID := testutil.ProjID()
	if projectID == "" {
		t.Skip("Need project ID. See CONTRIBUTING.md for details.")
	}
	ctx := context.Background()

	// Record.
	initial := time.Now()
	ibytes, err := json.Marshal(initial)
	if err != nil {
		t.Fatal(err)
	}
	rec, err := httpreplay.NewRecorder(replayFilename, ibytes)
	if err != nil {
		t.Fatal(err)
	}
	hc, err := rec.Client(ctx, option.WithTokenSource(
		testutil.TokenSource(ctx, storage.ScopeFullControl)))
	if err != nil {
		t.Fatal(err)
	}
	wanta, wantc := run(t, hc)
	testReadCRC(t, hc, "recording")
	if err := rec.Close(); err != nil {
		t.Fatalf("rec.Close: %v", err)
	}

	// Replay.
	rep, err := httpreplay.NewReplayer(replayFilename)
	if err != nil {
		t.Fatal(err)
	}
	defer rep.Close()
	hc, err = rep.Client(ctx)
	if err != nil {
		t.Fatal(err)
	}
	gota, gotc := run(t, hc)
	testReadCRC(t, hc, "replaying")

	if diff := testutil.Diff(gota, wanta); diff != "" {
		t.Error(diff)
	}
	if !bytes.Equal(gotc, wantc) {
		t.Errorf("got %q, want %q", gotc, wantc)
	}
	var gotInitial time.Time
	if err := json.Unmarshal(rep.Initial(), &gotInitial); err != nil {
		t.Fatal(err)
	}
	if !gotInitial.Equal(initial) {
		t.Errorf("initial: got %v, want %v", gotInitial, initial)
	}
}

// TODO(jba): test errors

func run(t *testing.T, hc *http.Client) (*storage.BucketAttrs, []byte) {
	ctx := context.Background()
	client, err := storage.NewClient(ctx, option.WithHTTPClient(hc))
	if err != nil {
		t.Fatal(err)
	}
	defer client.Close()
	b := client.Bucket(testutil.ProjID())
	attrs, err := b.Attrs(ctx)
	if err != nil {
		t.Fatal(err)
	}
	obj := b.Object("replay-test")
	w := obj.NewWriter(ctx)
	data := []byte{150, 151, 152}
	if _, err := w.Write(data); err != nil {
		t.Fatal(err)
	}
	if err := w.Close(); err != nil {
		t.Fatal(err)
	}

	r, err := obj.NewReader(ctx)
	if err != nil {
		t.Fatal(err)
	}
	defer r.Close()
	contents, err := ioutil.ReadAll(r)
	if err != nil {
		t.Fatal(err)
	}

	return attrs, contents
}

func testReadCRC(t *testing.T, hc *http.Client, mode string) {
	const (
		// This is an uncompressed file.
		// See https://cloud.google.com/storage/docs/public-datasets/landsat
		uncompressedBucket = "gcp-public-data-landsat"
		uncompressedObject = "LC08/PRE/044/034/LC80440342016259LGN00/LC80440342016259LGN00_MTL.txt"

		gzippedBucket = "storage-library-test-bucket"
		gzippedObject = "gzipped-text.txt"
	)
	ctx := context.Background()
	client, err := storage.NewClient(ctx, option.WithHTTPClient(hc))
	if err != nil {
		t.Fatalf("%s: %v", mode, err)
	}
	defer client.Close()

	uncompressedObj := client.Bucket(uncompressedBucket).Object(uncompressedObject)
	gzippedObj := client.Bucket(gzippedBucket).Object(gzippedObject)

	for _, test := range []struct {
		desc           string
		obj            *storage.ObjectHandle
		offset, length int64
		readCompressed bool // don't decompress a gzipped file

		wantErr bool
		wantLen int // length of contents
	}{
		{
			desc:           "uncompressed, entire file",
			obj:            uncompressedObj,
			offset:         0,
			length:         -1,
			readCompressed: false,
			wantLen:        7903,
		},
		{
			desc:           "uncompressed, entire file, don't decompress",
			obj:            uncompressedObj,
			offset:         0,
			length:         -1,
			readCompressed: true,
			wantLen:        7903,
		},
		{
			desc:           "uncompressed, suffix",
			obj:            uncompressedObj,
			offset:         3,
			length:         -1,
			readCompressed: false,
			wantLen:        7900,
		},
		{
			desc:           "uncompressed, prefix",
			obj:            uncompressedObj,
			offset:         0,
			length:         18,
			readCompressed: false,
			wantLen:        18,
		},
		{
			// When a gzipped file is unzipped by GCS, we can't verify the checksum
			// because it was computed against the zipped contents. There is no
			// header that indicates that a gzipped file is being served unzipped.
			// But our CRC check only happens if there is a Content-Length header,
			// and that header is absent for this read.
			desc:           "compressed, entire file, server unzips",
			obj:            gzippedObj,
			offset:         0,
			length:         -1,
			readCompressed: false,
			wantLen:        11,
		},
		{
			// When we read a gzipped file uncompressed, it's like reading a regular file:
			// the served content and the CRC match.
			desc:           "compressed, entire file, read compressed",
			obj:            gzippedObj,
			offset:         0,
			length:         -1,
			readCompressed: true,
			wantLen:        31,
		},
		{
			desc:           "compressed, partial, read compressed",
			obj:            gzippedObj,
			offset:         1,
			length:         8,
			readCompressed: true,
			wantLen:        8,
		},
		{
			desc:    "uncompressed, HEAD",
			obj:     uncompressedObj,
			offset:  0,
			length:  0,
			wantLen: 0,
		},
		{
			desc:    "compressed, HEAD",
			obj:     gzippedObj,
			offset:  0,
			length:  0,
			wantLen: 0,
		},
	} {
		obj := test.obj.ReadCompressed(test.readCompressed)
		r, err := obj.NewRangeReader(ctx, test.offset, test.length)
		if err != nil {
			if test.wantErr {
				continue
			}
			t.Errorf("%s: %s: %v", mode, test.desc, err)
			continue
		}
		data, err := ioutil.ReadAll(r)
		_ = r.Close()
		if err != nil {
			t.Errorf("%s: %s: %v", mode, test.desc, err)
			continue
		}
		if got, want := len(data), test.wantLen; got != want {
			t.Errorf("%s: %s: len: got %d, want %d", mode, test.desc, got, want)
		}
	}
}

func TestRemoveAndClear(t *testing.T) {
	// Disable logging for this test, since it generates a lot.
	log.SetOutput(ioutil.Discard)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		fmt.Fprintln(w, "LGTM")
	}))
	defer srv.Close()

	replayFilename := tempFilename(t, "TestRemoveAndClear*.replay")
	defer os.Remove(replayFilename)

	ctx := context.Background()
	// Record
	rec, err := httpreplay.NewRecorder(replayFilename, nil)
	if err != nil {
		t.Fatal(err)
	}
	rec.ClearHeaders("Clear")
	rec.RemoveRequestHeaders("Rem*")
	rec.ClearQueryParams("c")
	rec.RemoveQueryParams("r")
	hc, err := rec.Client(ctx, option.WithoutAuthentication())
	if err != nil {
		t.Fatal(err)
	}
	query := "k=1&r=2&c=3"
	req, err := http.NewRequest("GET", srv.URL+"?"+query, nil)
	if err != nil {
		t.Fatal(err)
	}
	headers := map[string]string{"Keep": "ok", "Clear": "secret", "Remove": "bye"}
	for k, v := range headers {
		req.Header.Set(k, v)
	}
	if _, err := hc.Do(req); err != nil {
		t.Fatal(err)
	}
	if err := rec.Close(); err != nil {
		t.Fatal(err)
	}

	// Replay
	// For both headers and query param:
	// - k or Keep must be present and identical
	// - c or Clear must be present, but can be different
	// - r or Remove can be anything
	for _, test := range []struct {
		query       string
		headers     map[string]string
		wantSuccess bool
	}{
		{query, headers, true}, // same query string and headers
		{query,
			map[string]string{"Keep": "oops", "Clear": "secret", "Remove": "bye"},
			false, // different Keep
		},
		{query, map[string]string{}, false},                               // missing Keep and Clear
		{query, map[string]string{"Keep": "ok"}, false},                   // missing Clear
		{query, map[string]string{"Keep": "ok", "Clear": "secret"}, true}, // missing Remove is OK
		{
			query,
			map[string]string{"Keep": "ok", "Clear": "secret", "Remove": "whatev"},
			true,
		}, // different Remove is OK
		{query, map[string]string{"Keep": "ok", "Clear": "diff"}, true}, // different Clear is OK
		{"", headers, false},            // no query string
		{"k=x&r=2&c=3", headers, false}, // different k
		{"r=2", headers, false},         // missing k and c
		{"k=1&r=2", headers, false},     // missing c
		{"k=1&c=3", headers, true},      // missing r is OK
		{"k=1&r=x&c=3", headers, true},  // different r is OK,
		{"k=1&r=2&c=x", headers, true},  // different clear is OK
	} {
		rep, err := httpreplay.NewReplayer(replayFilename)
		if err != nil {
			t.Fatal(err)
		}
		hc, err = rep.Client(ctx)
		if err != nil {
			t.Fatal(err)
		}
		url := srv.URL
		if test.query != "" {
			url += "?" + test.query
		}
		req, err = http.NewRequest("GET", url, nil)
		if err != nil {
			t.Fatal(err)
		}
		for k, v := range test.headers {
			req.Header.Set(k, v)
		}
		resp, err := hc.Do(req)
		if err != nil {
			t.Fatal(err)
		}
		rep.Close()
		if (resp.StatusCode == 200) != test.wantSuccess {
			t.Errorf("%q, %v: got %d, wanted success=%t",
				test.query, test.headers, resp.StatusCode, test.wantSuccess)
		}
	}
}

func tempFilename(t *testing.T, pattern string) string {
	f, err := ioutil.TempFile("", pattern)
	if err != nil {
		t.Fatal(err)
	}
	filename := f.Name()
	if err := f.Close(); err != nil {
		t.Fatal(err)
	}
	return filename
}
<|endoftext|>


<commit_before>package reverseproxy

import (
	"crypto/tls"
	"fmt"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strconv"
	"strings"
	"sync"

	"github.com/gorilla/websocket"
	"github.com/koding/kite"
	"github.com/koding/kite/config"
	"github.com/koding/websocketproxy"
)

const (
	Version = "0.0.1"
	Name    = "proxy"
)

type Proxy struct {
	Kite *kite.Kite

	listener  net.Listener
	TLSConfig *tls.Config

	readyC chan bool // To signal when kite is ready to accept connections
	closeC chan bool // To signal when kite is closed with Close()

	// Holds registered kites. Keys are kite IDs.
	kites   map[string]*url.URL
	kitesMu sync.Mutex

	// muxer for proxy
	mux            *http.ServeMux
	websocketProxy http.Handler
	httpProxy      http.Handler

	// Proxy properties used to give urls and bind the listener
	Scheme     string
	PublicHost string // If given it must match the domain in certificate.
	Port       int
}

func New(conf *config.Config) *Proxy {
	k := kite.New(Name, Version)
	k.Config = conf

	// Listen on 3999 by default

	p := &Proxy{
		Kite:   k,
		kites:  make(map[string]*url.URL),
		readyC: make(chan bool),
		closeC: make(chan bool),
		mux:    http.NewServeMux(),
		Port:   k.Config.Port,
	}

	// third part kites are going to use this to register themself to
	// proxy-kite and get a proxy url, which they use for register to kontrol.
	p.Kite.HandleFunc("register", p.handleRegister)

	// create our websocketproxy http.handler

	p.websocketProxy = &websocketproxy.WebsocketProxy{
		Backend: p.backend,
		Upgrader: &websocket.Upgrader{
			ReadBufferSize:  4096,
			WriteBufferSize: 4096,
			CheckOrigin: func(r *http.Request) bool {
				// TODO: change this to publicdomain and also kites should add them to
				return true
			},
		},
	}

	p.httpProxy = &httputil.ReverseProxy{
		Director: p.director,
	}

	p.mux.Handle("/", k)
	p.mux.Handle("/proxy/", p)

	// OnDisconnect is called whenever a kite is disconnected from us.
	k.OnDisconnect(func(r *kite.Client) {
		k.Log.Info("Removing kite Id '%s' from proxy. It's disconnected", r.Kite.ID)
		delete(p.kites, r.Kite.ID)
	})

	return p
}

// ServeHTTP implements the http.Handler interface.
func (p *Proxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	if isWebsocket(req) {
		p.websocketProxy.ServeHTTP(rw, req)
		return
	}

	p.httpProxy.ServeHTTP(rw, req)
}

// isWebsocket checks wether the incoming request is a part of websocket
// handshake
func isWebsocket(req *http.Request) bool {
	if strings.ToLower(req.Header.Get("Upgrade")) != "websocket" ||
		!strings.Contains(strings.ToLower(req.Header.Get("Connection")), "upgrade") {
		return false
	}
	return true
}

func (p *Proxy) CloseNotify() chan bool {
	return p.closeC
}

func (p *Proxy) ReadyNotify() chan bool {
	return p.readyC
}

func (p *Proxy) handleRegister(r *kite.Request) (interface{}, error) {
	kiteUrl, err := url.Parse(r.Args.One().MustString())
	if err != nil {
		return nil, err
	}

	p.kites[r.Client.ID] = kiteUrl

	proxyURL := url.URL{
		Scheme:   p.Scheme,
		Host:     p.PublicHost + ":" + strconv.Itoa(p.Port),
		Path:     "proxy",
		RawQuery: "kiteId=" + r.Client.ID,
	}

	s := proxyURL.String()
	p.Kite.Log.Info("Registering kite with url: '%s'. Can be reached now with: '%s'", kiteUrl, s)

	return s, nil
}

func (p *Proxy) director(req *http.Request) {
	u := p.backend(req)
	req.URL.Scheme = u.Scheme
	req.URL.Host = u.Host
	req.URL.Path = u.Path
}

func (p *Proxy) backend(req *http.Request) *url.URL {
	kiteId := req.URL.Query().Get("kiteId")

	p.kitesMu.Lock()
	defer p.kitesMu.Unlock()

	backendURL, ok := p.kites[kiteId]
	if !ok {
		p.Kite.Log.Error("kite for id '%s' is not found: %s", kiteId, req.URL.String())
		return nil
	}

	// change "http" with "ws" because websocket procol expects a ws or wss as
	// scheme
	if err := replaceSchemeWithWS(backendURL); err != nil {
		return nil
	}

	// change now the path for the backend kite. Kite register itself with
	// something like "localhost:7777/kite" however we are going to
	// dial/connect to a sockjs server and there is no sessionId/serverId in
	// the path. This causes problem because the SockJS serve can't parse it.
	// Therefore we as an intermediate client are getting the path as (ommited the query):
	// "/proxy/795/kite-fba0954a-07c7-4d34-4215-6a88733cf65c-OjLnvABL/websocket"
	// which will be converted to
	// "localhost:7777/kite/795/kite-fba0954a-07c7-4d34-4215-6a88733cf65c-OjLnvABL/websocket"
	backendURL.Path += strings.TrimPrefix(req.URL.Path, "/proxy")

	// also change the Origin to the client's host name, like as if someone
	// with the same backendUrl is trying to connect to the kite. Otherwise
	// will get an "Origin not allowed"
	req.Header.Set("Origin", "http://"+backendURL.Host)

	p.Kite.Log.Info("Returning backend url: '%s' for kiteId: %s", backendURL.String(), kiteId)
	return backendURL
}

// TODO: put this into a util package, is used by others too
func replaceSchemeWithWS(u *url.URL) error {
	switch u.Scheme {
	case "http":
		u.Scheme = "ws"
	case "https":
		u.Scheme = "wss"
	default:
		return fmt.Errorf("invalid scheme in url: %s", u.Scheme)
	}
	return nil
}

// ListenAndServe listens on the TCP network address addr and then calls Serve
// with handler to handle requests on incoming connections.
func (p *Proxy) ListenAndServe() error {
	var err error
	p.listener, err = net.Listen("tcp4",
		net.JoinHostPort(p.Kite.Config.IP, strconv.Itoa(p.Kite.Config.Port)))
	if err != nil {
		return err
	}
	p.Kite.Log.Info("Listening on: %s", p.listener.Addr().String())

	close(p.readyC)

	server := http.Server{
		Handler: p.mux,
	}

	defer close(p.closeC)
	return server.Serve(p.listener)
}

func (p *Proxy) ListenAndServeTLS(certFile, keyFile string) error {
	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		p.Kite.Log.Fatal("Could not load cert/key files: %s", err.Error())
	}

	tlsConfig := &tls.Config{
		Certificates: []tls.Certificate{cert},
	}

	p.listener, err = net.Listen("tcp",
		net.JoinHostPort(p.Kite.Config.IP, strconv.Itoa(p.Kite.Config.Port)))
	if err != nil {
		p.Kite.Log.Fatal(err.Error())
	}
	p.Kite.Log.Info("Listening on: %s", p.listener.Addr().String())

	// now we are ready
	close(p.readyC)

	p.listener = tls.NewListener(p.listener, tlsConfig)

	server := &http.Server{
		Handler:   p.mux,
		TLSConfig: tlsConfig,
	}

	defer close(p.closeC)
	return server.Serve(p.listener)
}

func (p *Proxy) Run() {
	p.ListenAndServe()
}
<commit_msg>reverseproxy: check for nilness<commit_after>package reverseproxy

import (
	"crypto/tls"
	"fmt"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strconv"
	"strings"
	"sync"

	"github.com/gorilla/websocket"
	"github.com/koding/kite"
	"github.com/koding/kite/config"
	"github.com/koding/websocketproxy"
)

const (
	Version = "0.0.1"
	Name    = "proxy"
)

type Proxy struct {
	Kite *kite.Kite

	listener  net.Listener
	TLSConfig *tls.Config

	readyC chan bool // To signal when kite is ready to accept connections
	closeC chan bool // To signal when kite is closed with Close()

	// Holds registered kites. Keys are kite IDs.
	kites   map[string]*url.URL
	kitesMu sync.Mutex

	// muxer for proxy
	mux            *http.ServeMux
	websocketProxy http.Handler
	httpProxy      http.Handler

	// Proxy properties used to give urls and bind the listener
	Scheme     string
	PublicHost string // If given it must match the domain in certificate.
	Port       int
}

func New(conf *config.Config) *Proxy {
	k := kite.New(Name, Version)
	k.Config = conf

	// Listen on 3999 by default

	p := &Proxy{
		Kite:   k,
		kites:  make(map[string]*url.URL),
		readyC: make(chan bool),
		closeC: make(chan bool),
		mux:    http.NewServeMux(),
		Port:   k.Config.Port,
	}

	// third part kites are going to use this to register themself to
	// proxy-kite and get a proxy url, which they use for register to kontrol.
	p.Kite.HandleFunc("register", p.handleRegister)

	// create our websocketproxy http.handler

	p.websocketProxy = &websocketproxy.WebsocketProxy{
		Backend: p.backend,
		Upgrader: &websocket.Upgrader{
			ReadBufferSize:  4096,
			WriteBufferSize: 4096,
			CheckOrigin: func(r *http.Request) bool {
				// TODO: change this to publicdomain and also kites should add them to
				return true
			},
		},
	}

	p.httpProxy = &httputil.ReverseProxy{
		Director: p.director,
	}

	p.mux.Handle("/", k)
	p.mux.Handle("/proxy/", p)

	// OnDisconnect is called whenever a kite is disconnected from us.
	k.OnDisconnect(func(r *kite.Client) {
		k.Log.Info("Removing kite Id '%s' from proxy. It's disconnected", r.Kite.ID)
		delete(p.kites, r.Kite.ID)
	})

	return p
}

// ServeHTTP implements the http.Handler interface.
func (p *Proxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	if isWebsocket(req) {
		p.websocketProxy.ServeHTTP(rw, req)
		return
	}

	p.httpProxy.ServeHTTP(rw, req)
}

// isWebsocket checks wether the incoming request is a part of websocket
// handshake
func isWebsocket(req *http.Request) bool {
	if strings.ToLower(req.Header.Get("Upgrade")) != "websocket" ||
		!strings.Contains(strings.ToLower(req.Header.Get("Connection")), "upgrade") {
		return false
	}
	return true
}

func (p *Proxy) CloseNotify() chan bool {
	return p.closeC
}

func (p *Proxy) ReadyNotify() chan bool {
	return p.readyC
}

func (p *Proxy) handleRegister(r *kite.Request) (interface{}, error) {
	kiteUrl, err := url.Parse(r.Args.One().MustString())
	if err != nil {
		return nil, err
	}

	p.kites[r.Client.ID] = kiteUrl

	proxyURL := url.URL{
		Scheme:   p.Scheme,
		Host:     p.PublicHost + ":" + strconv.Itoa(p.Port),
		Path:     "proxy",
		RawQuery: "kiteId=" + r.Client.ID,
	}

	s := proxyURL.String()
	p.Kite.Log.Info("Registering kite with url: '%s'. Can be reached now with: '%s'", kiteUrl, s)

	return s, nil
}

func (p *Proxy) director(req *http.Request) {
	u := p.backend(req)
	if u == nil {
		return
	}

	req.URL.Scheme = u.Scheme
	req.URL.Host = u.Host
	req.URL.Path = u.Path
}

func (p *Proxy) backend(req *http.Request) *url.URL {
	kiteId := req.URL.Query().Get("kiteId")

	p.kitesMu.Lock()
	defer p.kitesMu.Unlock()

	backendURL, ok := p.kites[kiteId]
	if !ok {
		p.Kite.Log.Error("kite for id '%s' is not found: %s", kiteId, req.URL.String())
		return nil
	}

	// change "http" with "ws" because websocket procol expects a ws or wss as
	// scheme
	if err := replaceSchemeWithWS(backendURL); err != nil {
		return nil
	}

	// change now the path for the backend kite. Kite register itself with
	// something like "localhost:7777/kite" however we are going to
	// dial/connect to a sockjs server and there is no sessionId/serverId in
	// the path. This causes problem because the SockJS serve can't parse it.
	// Therefore we as an intermediate client are getting the path as (ommited the query):
	// "/proxy/795/kite-fba0954a-07c7-4d34-4215-6a88733cf65c-OjLnvABL/websocket"
	// which will be converted to
	// "localhost:7777/kite/795/kite-fba0954a-07c7-4d34-4215-6a88733cf65c-OjLnvABL/websocket"
	backendURL.Path += strings.TrimPrefix(req.URL.Path, "/proxy")

	// also change the Origin to the client's host name, like as if someone
	// with the same backendUrl is trying to connect to the kite. Otherwise
	// will get an "Origin not allowed"
	req.Header.Set("Origin", "http://"+backendURL.Host)

	p.Kite.Log.Info("Returning backend url: '%s' for kiteId: %s", backendURL.String(), kiteId)
	return backendURL
}

// TODO: put this into a util package, is used by others too
func replaceSchemeWithWS(u *url.URL) error {
	switch u.Scheme {
	case "http":
		u.Scheme = "ws"
	case "https":
		u.Scheme = "wss"
	default:
		return fmt.Errorf("invalid scheme in url: %s", u.Scheme)
	}
	return nil
}

// ListenAndServe listens on the TCP network address addr and then calls Serve
// with handler to handle requests on incoming connections.
func (p *Proxy) ListenAndServe() error {
	var err error
	p.listener, err = net.Listen("tcp4",
		net.JoinHostPort(p.Kite.Config.IP, strconv.Itoa(p.Kite.Config.Port)))
	if err != nil {
		return err
	}
	p.Kite.Log.Info("Listening on: %s", p.listener.Addr().String())

	close(p.readyC)

	server := http.Server{
		Handler: p.mux,
	}

	defer close(p.closeC)
	return server.Serve(p.listener)
}

func (p *Proxy) ListenAndServeTLS(certFile, keyFile string) error {
	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		p.Kite.Log.Fatal("Could not load cert/key files: %s", err.Error())
	}

	tlsConfig := &tls.Config{
		Certificates: []tls.Certificate{cert},
	}

	p.listener, err = net.Listen("tcp",
		net.JoinHostPort(p.Kite.Config.IP, strconv.Itoa(p.Kite.Config.Port)))
	if err != nil {
		p.Kite.Log.Fatal(err.Error())
	}
	p.Kite.Log.Info("Listening on: %s", p.listener.Addr().String())

	// now we are ready
	close(p.readyC)

	p.listener = tls.NewListener(p.listener, tlsConfig)

	server := &http.Server{
		Handler:   p.mux,
		TLSConfig: tlsConfig,
	}

	defer close(p.closeC)
	return server.Serve(p.listener)
}

func (p *Proxy) Run() {
	p.ListenAndServe()
}
<|endoftext|>


<commit_before>package platform

// Owner represents a resource owner
type Owner struct {
	ID ID
}
<commit_msg>add decode method to Owner<commit_after>package platform

import "encoding/hex"

// Owner represents a resource owner
type Owner struct {
	ID ID
}

// Decode parses b as a hex-encoded byte-slice-string.
func (o *Owner) Decode(b []byte) error {
	dst := make([]byte, hex.DecodedLen(len(b)))
	_, err := hex.Decode(dst, b)
	if err != nil {
		return err
	}
	o.ID = dst
	return nil
}
<|endoftext|>


<commit_before>package vips

/*
#cgo pkg-config: vips
#include "vips-bridge.h"
*/
import "C"

import (
	"fmt"
	"runtime"
	"sync"
	"unsafe"

	"github.com/DMarby/picsum-photos/logger"
)

// Image is a representation of the *C.VipsImage type
type Image *C.VipsImage

var (
	once sync.Once
	log  *logger.Logger
)

// Initialize libvips if it's not already started
func Initialize(logger *logger.Logger) error {
	var err error

	log = logger

	once.Do(func() {
		// vips_init needs to run on the main thread
		runtime.LockOSThread()
		defer runtime.UnlockOSThread()

		if C.VIPS_MAJOR_VERSION != 8 || C.VIPS_MINOR_VERSION < 6 {
			err = fmt.Errorf("unsupported libvips version")
		}

		errorCode := C.vips_init(C.CString("picsum-photos"))
		if errorCode != 0 {
			err = fmt.Errorf("unable to initialize vips: %v", catchVipsError())
			return
		}

		// Catch vips logging/warnings
		C.setup_logging()

		// Set concurrency to 1 so that each job only uses one thread
		C.vips_concurrency_set(1)
	})

	return err
}

// log_callback catches logs from libvips
//export log_callback
func log_callback(message *C.char) {
	log.Debug(C.GoString(message))
}

// Shutdown libvips
func Shutdown() {
	C.vips_shutdown()
}

// PrintDebugInfo prints libvips debug info to stdout
func PrintDebugInfo() {
	C.vips_object_print_all()
}

// catchVipsError returns the vips error buffer as an error
func catchVipsError() error {
	defer C.vips_error_clear()

	s := C.GoString(C.vips_error_buffer())
	return fmt.Errorf("%s", s)
}

// ResizeImage loads an image from a buffer and resizes it.
func ResizeImage(buffer []byte, width int, height int) (Image, error) {
	if len(buffer) == 0 {
		return nil, fmt.Errorf("empty buffer")
	}

	imageBuffer := unsafe.Pointer(&buffer[0])
	imageBufferSize := C.size_t(len(buffer))

	var image *C.VipsImage

	errCode := C.resize_image(imageBuffer, imageBufferSize, &image, C.int(width), C.int(height), C.VIPS_INTERESTING_CENTRE)

	// Prevent buffer from being garbage collected until after resize_image has been called
	runtime.KeepAlive(buffer)

	if errCode != 0 {
		return nil, fmt.Errorf("error processing image from buffer %s", catchVipsError())
	}

	return image, nil
}

// SaveToBuffer saves an image to a buffer
func SaveToBuffer(image Image) ([]byte, error) {
	defer UnrefImage(image)

	var bufferPointer unsafe.Pointer
	bufferLength := C.size_t(0)

	err := C.save_image_to_jpeg_buffer(image, &bufferPointer, &bufferLength)

	if err != 0 {
		return nil, fmt.Errorf("error saving to buffer %s", catchVipsError())
	}

	buffer := C.GoBytes(bufferPointer, C.int(bufferLength))

	C.g_free(C.gpointer(bufferPointer))

	return buffer, nil
}

// Grayscale converts an image to grayscale
func Grayscale(image Image) (Image, error) {
	defer UnrefImage(image)

	var result *C.VipsImage

	err := C.change_colorspace(image, &result, C.VIPS_INTERPRETATION_B_W)

	if err != 0 {
		return nil, fmt.Errorf("error changing image colorspace %s", catchVipsError())
	}

	return result, nil
}

// Blur applies gaussian blur to an image
func Blur(image Image, blur int) (Image, error) {
	defer UnrefImage(image)

	var result *C.VipsImage

	err := C.blur_image(image, &result, C.double(blur))

	if err != 0 {
		return nil, fmt.Errorf("error applying blur to image %s", catchVipsError())
	}

	return result, nil
}

// UnrefImage unrefs an image object
func UnrefImage(image Image) {
	C.g_object_unref(C.gpointer(image))
}

// NewEmptyImage returns an empty image object
func NewEmptyImage() Image {
	return C.vips_image_new()
}
<commit_msg>Disable the vips cache<commit_after>package vips

/*
#cgo pkg-config: vips
#include "vips-bridge.h"
*/
import "C"

import (
	"fmt"
	"runtime"
	"sync"
	"unsafe"

	"github.com/DMarby/picsum-photos/logger"
)

// Image is a representation of the *C.VipsImage type
type Image *C.VipsImage

var (
	once sync.Once
	log  *logger.Logger
)

// Initialize libvips if it's not already started
func Initialize(logger *logger.Logger) error {
	var err error

	log = logger

	once.Do(func() {
		// vips_init needs to run on the main thread
		runtime.LockOSThread()
		defer runtime.UnlockOSThread()

		if C.VIPS_MAJOR_VERSION != 8 || C.VIPS_MINOR_VERSION < 6 {
			err = fmt.Errorf("unsupported libvips version")
		}

		errorCode := C.vips_init(C.CString("picsum-photos"))
		if errorCode != 0 {
			err = fmt.Errorf("unable to initialize vips: %v", catchVipsError())
			return
		}

		// Catch vips logging/warnings
		C.setup_logging()

		// Set concurrency to 1 so that each job only uses one thread
		C.vips_concurrency_set(1)

		// Disable the cache
		C.vips_cache_set_max_mem(0)
		C.vips_cache_set_max(0)
	})

	return err
}

// log_callback catches logs from libvips
//export log_callback
func log_callback(message *C.char) {
	log.Debug(C.GoString(message))
}

// Shutdown libvips
func Shutdown() {
	C.vips_shutdown()
}

// PrintDebugInfo prints libvips debug info to stdout
func PrintDebugInfo() {
	C.vips_object_print_all()
}

// catchVipsError returns the vips error buffer as an error
func catchVipsError() error {
	defer C.vips_error_clear()

	s := C.GoString(C.vips_error_buffer())
	return fmt.Errorf("%s", s)
}

// ResizeImage loads an image from a buffer and resizes it.
func ResizeImage(buffer []byte, width int, height int) (Image, error) {
	if len(buffer) == 0 {
		return nil, fmt.Errorf("empty buffer")
	}

	imageBuffer := unsafe.Pointer(&buffer[0])
	imageBufferSize := C.size_t(len(buffer))

	var image *C.VipsImage

	errCode := C.resize_image(imageBuffer, imageBufferSize, &image, C.int(width), C.int(height), C.VIPS_INTERESTING_CENTRE)

	// Prevent buffer from being garbage collected until after resize_image has been called
	runtime.KeepAlive(buffer)

	if errCode != 0 {
		return nil, fmt.Errorf("error processing image from buffer %s", catchVipsError())
	}

	return image, nil
}

// SaveToBuffer saves an image to a buffer
func SaveToBuffer(image Image) ([]byte, error) {
	defer UnrefImage(image)

	var bufferPointer unsafe.Pointer
	bufferLength := C.size_t(0)

	err := C.save_image_to_jpeg_buffer(image, &bufferPointer, &bufferLength)

	if err != 0 {
		return nil, fmt.Errorf("error saving to buffer %s", catchVipsError())
	}

	buffer := C.GoBytes(bufferPointer, C.int(bufferLength))

	C.g_free(C.gpointer(bufferPointer))

	return buffer, nil
}

// Grayscale converts an image to grayscale
func Grayscale(image Image) (Image, error) {
	defer UnrefImage(image)

	var result *C.VipsImage

	err := C.change_colorspace(image, &result, C.VIPS_INTERPRETATION_B_W)

	if err != 0 {
		return nil, fmt.Errorf("error changing image colorspace %s", catchVipsError())
	}

	return result, nil
}

// Blur applies gaussian blur to an image
func Blur(image Image, blur int) (Image, error) {
	defer UnrefImage(image)

	var result *C.VipsImage

	err := C.blur_image(image, &result, C.double(blur))

	if err != 0 {
		return nil, fmt.Errorf("error applying blur to image %s", catchVipsError())
	}

	return result, nil
}

// UnrefImage unrefs an image object
func UnrefImage(image Image) {
	C.g_object_unref(C.gpointer(image))
}

// NewEmptyImage returns an empty image object
func NewEmptyImage() Image {
	return C.vips_image_new()
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2016 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 server

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/abbot/go-http-auth"
	"github.com/skydive-project/skydive/api/types"
	"github.com/skydive-project/skydive/flow"
	ge "github.com/skydive-project/skydive/gremlin/traversal"
	shttp "github.com/skydive-project/skydive/http"
	"github.com/skydive-project/skydive/logging"
	"github.com/skydive-project/skydive/rbac"
	"github.com/skydive-project/skydive/topology/graph"
	"github.com/skydive-project/skydive/topology/graph/traversal"
	"github.com/skydive-project/skydive/validator"
)

// TopologyAPI exposes the topology query API
type TopologyAPI struct {
	graph         *graph.Graph
	gremlinParser *traversal.GremlinTraversalParser
}

func shortID(s graph.Identifier) graph.Identifier {
	if len(s) > 8 {
		return s[:8]
	}
	return s
}

func (t *TopologyAPI) graphToDot(w http.ResponseWriter, g *graph.Graph) {
	g.RLock()
	defer g.RUnlock()

	w.Write([]byte("digraph g {\n"))

	nodeMap := make(map[graph.Identifier]*graph.Node)
	for _, n := range g.GetNodes(nil) {
		nodeMap[n.ID] = n
		name, _ := n.GetFieldString("Name")
		title := fmt.Sprintf("%s-%s", name, shortID(n.ID))
		label := title
		for k, v := range n.Metadata() {
			switch k {
			case "Type", "IfIndex", "State", "TID":
				label += fmt.Sprintf("\\n%s = %v", k, v)
			}
		}
		w.Write([]byte(fmt.Sprintf("\"%s\" [label=\"%s\"]\n", title, label)))
	}

	for _, e := range g.GetEdges(nil) {
		parent := nodeMap[e.GetParent()]
		child := nodeMap[e.GetChild()]
		if parent == nil || child == nil {
			continue
		}

		childName, _ := child.GetFieldString("Name")
		parentName, _ := parent.GetFieldString("Name")
		relationType, _ := e.GetFieldString("RelationType")
		linkLabel, linkType, direction := "", "->", "forward"
		switch relationType {
		case "":
		case "layer2":
			direction = "both"
			fallthrough
		default:
			linkLabel = fmt.Sprintf(" [label=%s,dir=%s]\n", relationType, direction)
		}
		link := fmt.Sprintf("\"%s-%s\" %s \"%s-%s\"%s", parentName, shortID(parent.ID), linkType, childName, shortID(child.ID), linkLabel)
		w.Write([]byte(link))
	}

	w.Write([]byte("}"))
}

func (t *TopologyAPI) topologyIndex(w http.ResponseWriter, r *auth.AuthenticatedRequest) {
	if !rbac.Enforce(r.Username, "topology", "read") {
		w.WriteHeader(http.StatusMethodNotAllowed)
		return
	}

	t.graph.RLock()
	defer t.graph.RUnlock()

	w.WriteHeader(http.StatusOK)
	if strings.Contains(r.Header.Get("Accept"), "vnd.graphviz") {
		w.Header().Set("Content-Type", "text/vnd.graphviz; charset=UTF-8")
		t.graphToDot(w, t.graph)
	} else {
		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		if err := json.NewEncoder(w).Encode(t.graph); err != nil {
			logging.GetLogger().Warningf("Error while writing response: %s", err)
		}
	}
}

func (t *TopologyAPI) topologySearch(w http.ResponseWriter, r *auth.AuthenticatedRequest) {
	if !rbac.Enforce(r.Username, "topology", "read") {
		w.WriteHeader(http.StatusMethodNotAllowed)
		return
	}

	resource := types.TopologyParam{}
	data, _ := ioutil.ReadAll(r.Body)
	if len(data) != 0 {
		if err := json.Unmarshal(data, &resource); err != nil {
			writeError(w, http.StatusBadRequest, err)
			return
		}
		if err := validator.Validate(resource); err != nil {
			writeError(w, http.StatusBadRequest, err)
			return
		}
	}

	if resource.GremlinQuery == "" {
		w.WriteHeader(http.StatusNoContent)
		return
	}

	ts, err := t.gremlinParser.Parse(strings.NewReader(resource.GremlinQuery))
	if err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}

	res, err := ts.Exec(t.graph, true)
	if err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}

	if strings.Contains(r.Header.Get("Accept"), "vnd.graphviz") {
		if graphTraversal, ok := res.(*traversal.GraphTraversal); ok {
			w.Header().Set("Content-Type", "text/vnd.graphviz; charset=UTF-8")
			w.WriteHeader(http.StatusOK)
			t.graphToDot(w, graphTraversal.Graph)
		} else {
			writeError(w, http.StatusNotAcceptable, errors.New("Only graph can be outputted as dot"))
		}
	} else if strings.Contains(r.Header.Get("Accept"), "vnd.tcpdump.pcap") {
		if rawPacketsTraversal, ok := res.(*ge.RawPacketsTraversalStep); ok {
			values := rawPacketsTraversal.Values()
			if len(values) == 0 {
				writeError(w, http.StatusNotFound, errors.New("No raw packet found, please check your Gremlin request and the time context"))
			} else {
				w.Header().Set("Content-Type", "application/vnd.tcpdump.pcap; charset=UTF-8")
				w.WriteHeader(http.StatusOK)

				pw := flow.NewPcapWriter(w)
				for _, pf := range values {
					m := pf.(map[string]*flow.RawPackets)
					for _, fr := range m {
						if err = pw.WriteRawPackets(fr); err != nil {
							writeError(w, http.StatusNotAcceptable, errors.New(err.Error()))
						}
					}
				}
			}
		} else {
			writeError(w, http.StatusNotAcceptable, errors.New("Only RawPackets step result can be outputted as pcap"))
		}
	} else {
		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(http.StatusOK)
		if err := json.NewEncoder(w).Encode(res); err != nil {
			logging.GetLogger().Warningf("Error while writing response: %s", err)
		}
	}
}

func (t *TopologyAPI) registerEndpoints(r *shttp.Server, authBackend shttp.AuthenticationBackend) {
	routes := []shttp.Route{
		{
			Name:        "TopologiesIndex",
			Method:      "GET",
			Path:        "/api/topology",
			HandlerFunc: t.topologyIndex,
		},
		{
			Name:        "TopologiesSearch",
			Method:      "POST",
			Path:        "/api/topology",
			HandlerFunc: t.topologySearch,
		},
	}

	r.RegisterRoutes(routes, authBackend)
}

// RegisterTopologyAPI registers a new topology query API
func RegisterTopologyAPI(r *shttp.Server, g *graph.Graph, parser *traversal.GremlinTraversalParser, authBackend shttp.AuthenticationBackend) {
	t := &TopologyAPI{
		gremlinParser: parser,
		graph:         g,
	}

	t.registerEndpoints(r, authBackend)
}
<commit_msg>api: display IP addresses of nodes in dot output<commit_after>/*
 * Copyright (C) 2016 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 server

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/abbot/go-http-auth"
	"github.com/skydive-project/skydive/api/types"
	"github.com/skydive-project/skydive/flow"
	ge "github.com/skydive-project/skydive/gremlin/traversal"
	shttp "github.com/skydive-project/skydive/http"
	"github.com/skydive-project/skydive/logging"
	"github.com/skydive-project/skydive/rbac"
	"github.com/skydive-project/skydive/topology/graph"
	"github.com/skydive-project/skydive/topology/graph/traversal"
	"github.com/skydive-project/skydive/validator"
)

// TopologyAPI exposes the topology query API
type TopologyAPI struct {
	graph         *graph.Graph
	gremlinParser *traversal.GremlinTraversalParser
}

func shortID(s graph.Identifier) graph.Identifier {
	if len(s) > 8 {
		return s[:8]
	}
	return s
}

func (t *TopologyAPI) graphToDot(w http.ResponseWriter, g *graph.Graph) {
	g.RLock()
	defer g.RUnlock()

	w.Write([]byte("digraph g {\n"))

	nodeMap := make(map[graph.Identifier]*graph.Node)
	for _, n := range g.GetNodes(nil) {
		nodeMap[n.ID] = n
		name, _ := n.GetFieldString("Name")
		title := fmt.Sprintf("%s-%s", name, shortID(n.ID))
		label := title
		for k, v := range n.Metadata() {
			switch k {
			case "Type", "IfIndex", "State", "TID", "IPV4", "IPV6":
				label += fmt.Sprintf("\\n%s = %v", k, v)
			}
		}
		w.Write([]byte(fmt.Sprintf("\"%s\" [label=\"%s\"]\n", title, label)))
	}

	for _, e := range g.GetEdges(nil) {
		parent := nodeMap[e.GetParent()]
		child := nodeMap[e.GetChild()]
		if parent == nil || child == nil {
			continue
		}

		childName, _ := child.GetFieldString("Name")
		parentName, _ := parent.GetFieldString("Name")
		relationType, _ := e.GetFieldString("RelationType")
		linkLabel, linkType, direction := "", "->", "forward"
		switch relationType {
		case "":
		case "layer2":
			direction = "both"
			fallthrough
		default:
			linkLabel = fmt.Sprintf(" [label=%s,dir=%s]\n", relationType, direction)
		}
		link := fmt.Sprintf("\"%s-%s\" %s \"%s-%s\"%s", parentName, shortID(parent.ID), linkType, childName, shortID(child.ID), linkLabel)
		w.Write([]byte(link))
	}

	w.Write([]byte("}"))
}

func (t *TopologyAPI) topologyIndex(w http.ResponseWriter, r *auth.AuthenticatedRequest) {
	if !rbac.Enforce(r.Username, "topology", "read") {
		w.WriteHeader(http.StatusMethodNotAllowed)
		return
	}

	t.graph.RLock()
	defer t.graph.RUnlock()

	w.WriteHeader(http.StatusOK)
	if strings.Contains(r.Header.Get("Accept"), "vnd.graphviz") {
		w.Header().Set("Content-Type", "text/vnd.graphviz; charset=UTF-8")
		t.graphToDot(w, t.graph)
	} else {
		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		if err := json.NewEncoder(w).Encode(t.graph); err != nil {
			logging.GetLogger().Warningf("Error while writing response: %s", err)
		}
	}
}

func (t *TopologyAPI) topologySearch(w http.ResponseWriter, r *auth.AuthenticatedRequest) {
	if !rbac.Enforce(r.Username, "topology", "read") {
		w.WriteHeader(http.StatusMethodNotAllowed)
		return
	}

	resource := types.TopologyParam{}
	data, _ := ioutil.ReadAll(r.Body)
	if len(data) != 0 {
		if err := json.Unmarshal(data, &resource); err != nil {
			writeError(w, http.StatusBadRequest, err)
			return
		}
		if err := validator.Validate(resource); err != nil {
			writeError(w, http.StatusBadRequest, err)
			return
		}
	}

	if resource.GremlinQuery == "" {
		w.WriteHeader(http.StatusNoContent)
		return
	}

	ts, err := t.gremlinParser.Parse(strings.NewReader(resource.GremlinQuery))
	if err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}

	res, err := ts.Exec(t.graph, true)
	if err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}

	if strings.Contains(r.Header.Get("Accept"), "vnd.graphviz") {
		if graphTraversal, ok := res.(*traversal.GraphTraversal); ok {
			w.Header().Set("Content-Type", "text/vnd.graphviz; charset=UTF-8")
			w.WriteHeader(http.StatusOK)
			t.graphToDot(w, graphTraversal.Graph)
		} else {
			writeError(w, http.StatusNotAcceptable, errors.New("Only graph can be outputted as dot"))
		}
	} else if strings.Contains(r.Header.Get("Accept"), "vnd.tcpdump.pcap") {
		if rawPacketsTraversal, ok := res.(*ge.RawPacketsTraversalStep); ok {
			values := rawPacketsTraversal.Values()
			if len(values) == 0 {
				writeError(w, http.StatusNotFound, errors.New("No raw packet found, please check your Gremlin request and the time context"))
			} else {
				w.Header().Set("Content-Type", "application/vnd.tcpdump.pcap; charset=UTF-8")
				w.WriteHeader(http.StatusOK)

				pw := flow.NewPcapWriter(w)
				for _, pf := range values {
					m := pf.(map[string]*flow.RawPackets)
					for _, fr := range m {
						if err = pw.WriteRawPackets(fr); err != nil {
							writeError(w, http.StatusNotAcceptable, errors.New(err.Error()))
						}
					}
				}
			}
		} else {
			writeError(w, http.StatusNotAcceptable, errors.New("Only RawPackets step result can be outputted as pcap"))
		}
	} else {
		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(http.StatusOK)
		if err := json.NewEncoder(w).Encode(res); err != nil {
			logging.GetLogger().Warningf("Error while writing response: %s", err)
		}
	}
}

func (t *TopologyAPI) registerEndpoints(r *shttp.Server, authBackend shttp.AuthenticationBackend) {
	routes := []shttp.Route{
		{
			Name:        "TopologiesIndex",
			Method:      "GET",
			Path:        "/api/topology",
			HandlerFunc: t.topologyIndex,
		},
		{
			Name:        "TopologiesSearch",
			Method:      "POST",
			Path:        "/api/topology",
			HandlerFunc: t.topologySearch,
		},
	}

	r.RegisterRoutes(routes, authBackend)
}

// RegisterTopologyAPI registers a new topology query API
func RegisterTopologyAPI(r *shttp.Server, g *graph.Graph, parser *traversal.GremlinTraversalParser, authBackend shttp.AuthenticationBackend) {
	t := &TopologyAPI{
		gremlinParser: parser,
		graph:         g,
	}

	t.registerEndpoints(r, authBackend)
}
<|endoftext|>


<commit_before>package registry

import (
	"errors"
	"path"
	"strings"

	etcdErr "github.com/coreos/fleet/third_party/github.com/coreos/etcd/error"
	"github.com/coreos/fleet/third_party/github.com/coreos/go-etcd/etcd"
	log "github.com/coreos/fleet/third_party/github.com/golang/glog"

	"github.com/coreos/fleet/event"
	"github.com/coreos/fleet/job"
	"github.com/coreos/fleet/unit"
)

const (
	jobPrefix = "job"
)

// GetAllJobs lists all Jobs known by the Registry
func (r *Registry) GetAllJobs() []job.Job {
	var jobs []job.Job

	key := path.Join(r.keyPrefix, jobPrefix)
	resp, err := r.etcd.Get(key, true, true)

	if err != nil {
		return jobs
	}

	for _, dir := range resp.Node.Nodes {
		for _, node := range dir.Nodes {
			if !strings.HasSuffix(node.Key, "object") {
				continue
			}

			j := r.getJobFromJSON(node.Value)
			if j == nil {
				continue
			}

			jobs = append(jobs, *j)
		}
	}

	return jobs
}

// GetJobTarget looks up where the given job is scheduled. If the job has
// been scheduled, the ID the target machine is returned. Otherwise, an
// empty string is returned.
func (r *Registry) GetJobTarget(jobName string) string {
	// Figure out to which Machine this Job is scheduled
	key := r.jobTargetAgentPath(jobName)
	resp, err := r.etcd.Get(key, false, true)
	if err != nil {
		return ""
	}

	return resp.Node.Value
}

func (r *Registry) ClearJobTarget(jobName, machID string) error {
	key := r.jobTargetAgentPath(jobName)
	_, err := r.etcd.CompareAndDelete(key, machID, 0)
	return err
}

// GetJob looks for a Job of the given name in the Registry. It returns a fully
// hydrated Job on success, or nil on any kind of failure.
func (r *Registry) GetJob(jobName string) (j *job.Job) {
	key := path.Join(r.keyPrefix, jobPrefix, jobName, "object")
	resp, err := r.etcd.Get(key, false, true)

	// Assume the error was KeyNotFound and return an empty data structure
	if err != nil {
		return nil
	}

	j = r.getJobFromJSON(resp.Node.Value)
	if j == nil {
		log.V(1).Infof("Error unmarshaling Job(%s): %v", jobName, err)
	}
	return
}

func (r *Registry) getJobFromJSON(val string) *job.Job {
	var jm jobModel
	if err := unmarshal(val, &jm); err != nil {
		return nil
	}

	return r.getJobFromModel(jm)
}

func (r *Registry) getJobFromModel(jm jobModel) *job.Job {
	var err error
	var unit *unit.Unit

	// New-style Jobs should have a populated UnitHash, and the contents of the Unit are stored separately in the Registry
	if !jm.UnitHash.Empty() {
		unit = r.getUnitByHash(jm.UnitHash)
		if unit == nil {
			log.Warningf("No Unit found in Registry for Job(%s)", jm.Name)
			return nil
		}
		if unit.Hash() != jm.UnitHash {
			log.Errorf("Unit Hash %s does not match expected %s for Job(%s)!", unit.Hash(), jm.UnitHash, jm.Name)
			return nil
		}
		log.V(2).Infof("Got Unit for Job(%s) from registry", jm.Name)
	} else {
		// Old-style Jobs had "Payloads" instead of Units, also stored separately in the Registry
		log.V(2).Infof("Legacy Job(%s) has no PayloadHash - looking for associated Payload", jm.Name)
		unit, err = r.getUnitFromLegacyPayload(jm.Name)
		if err != nil {
			log.Errorf("Error retrieving legacy payload for Job(%s)", jm.Name)
			return nil
		} else if unit == nil {
			log.Warningf("No Payload found in Registry for Job(%s)", jm.Name)
			return nil
		}
	}

	j := job.NewJob(jm.Name, *unit)

	j.UnitState = r.getUnitState(jm.Name)
	j.State = r.determineJobState(jm.Name)

	return j
}

// jobModel is used for serializing and deserializing Jobs stored in the Registry
type jobModel struct {
	Name     string
	UnitHash unit.Hash
}

// DestroyJob removes a Job object from the repository, along with any legacy
// associated Payload and SignatureSet. It does not yet remove underlying
// Units from the repository.
func (r *Registry) DestroyJob(jobName string) {
	key := path.Join(r.keyPrefix, jobPrefix, jobName)
	r.etcd.Delete(key, true)
	// TODO(jonboulle): add unit reference counting and actually destroying Units
	r.destroyLegacyPayload(jobName)
	r.destroySignatureSetOfJob(jobName)
}

// destroyLegacyPayload removes an old-style Payload from the registry
func (r *Registry) destroyLegacyPayload(payloadName string) {
	key := path.Join(r.keyPrefix, payloadPrefix, payloadName)
	r.etcd.Delete(key, false)
}

// CreateJob attempts to store a Job and its associated Unit in the registry
func (r *Registry) CreateJob(j *job.Job) (err error) {
	if err := r.storeOrGetUnit(j.Unit); err != nil {
		return err
	}

	key := path.Join(r.keyPrefix, jobPrefix, j.Name, "object")

	jm := jobModel{
		Name:     j.Name,
		UnitHash: j.UnitHash,
	}
	json, err := marshal(jm)
	if err != nil {
		return
	}

	_, err = r.etcd.Create(key, json, 0)
	if err != nil && err.(*etcd.EtcdError).ErrorCode == etcdErr.EcodeNodeExist {
		err = errors.New("job already exists")
	}

	return
}

func (r *Registry) GetJobTargetState(jobName string) *job.JobState {
	key := r.jobTargetStatePath(jobName)
	resp, err := r.etcd.Get(key, false, false)
	if err != nil {
		if err.(*etcd.EtcdError).ErrorCode != etcdErr.EcodeNodeExist {
			log.Errorf("Unable to determine target-state of Job(%s): %v", jobName, err)
		}
		return nil
	}

	return job.ParseJobState(resp.Node.Value)
}

func (r *Registry) SetJobTargetState(jobName string, state job.JobState) error {
	key := r.jobTargetStatePath(jobName)
	_, err := r.etcd.Set(key, string(state), 0)
	return err
}

func (es *EventStream) filterJobTargetStateChanges(resp *etcd.Response) *event.Event {
	if resp.Action != "set" {
		return nil
	}

	dir, baseName := path.Split(resp.Node.Key)
	if baseName != "target-state" {
		return nil
	}

	dir = strings.TrimSuffix(dir, "/")
	jobName := path.Base(dir)

	ts := job.ParseJobState(resp.Node.Value)
	if ts == nil {
		return nil
	}

	cs := es.registry.determineJobState(jobName)
	if *cs == *ts {
		return nil
	}

	var cType string
	switch *cs {
	case job.JobStateInactive:
		cType = "CommandLoadJob"
	case job.JobStateLoaded:
		if *ts == job.JobStateInactive {
			cType = "CommandUnloadJob"
		} else if *ts == job.JobStateLaunched {
			cType = "CommandStartJob"
		}
	case job.JobStateLaunched:
		if *ts == job.JobStateLoaded {
			cType = "CommandStopJob"
		} else if *ts == job.JobStateInactive {
			cType = "CommandUnloadJob"
		}
	}

	if cType == "" {
		return nil
	}

	agent := es.registry.GetJobTarget(jobName)
	return &event.Event{cType, jobName, agent}
}

func (r *Registry) ScheduleJob(jobName string, machID string) error {
	key := r.jobTargetAgentPath(jobName)
	_, err := r.etcd.Create(key, machID, 0)
	return err
}

func (r *Registry) LockJob(jobName, context string) *TimedResourceMutex {
	return r.lockResource("job", jobName, context)
}

func filterEventJobScheduled(resp *etcd.Response) *event.Event {
	if resp.Action != "create" {
		return nil
	}

	dir, baseName := path.Split(resp.Node.Key)
	if baseName != "target" {
		return nil
	}

	dir = strings.TrimSuffix(dir, "/")
	dir, jobName := path.Split(dir)

	dir = strings.TrimSuffix(dir, "/")
	dir, prefixName := path.Split(dir)

	if prefixName != jobPrefix {
		return nil
	}

	return &event.Event{"EventJobScheduled", jobName, resp.Node.Value}
}

func filterEventJobUnscheduled(resp *etcd.Response) *event.Event {
	if resp.Action != "delete" && resp.Action != "compareAndDelete" {
		return nil
	}

	dir, baseName := path.Split(resp.Node.Key)
	if baseName != "target" {
		return nil
	}

	dir = strings.TrimSuffix(dir, "/")
	dir, jobName := path.Split(dir)

	dir = strings.TrimSuffix(dir, "/")
	dir, prefixName := path.Split(dir)

	if prefixName != jobPrefix {
		return nil
	}

	if resp.PrevNode == nil {
		return nil
	}

	return &event.Event{"EventJobUnscheduled", jobName, resp.PrevNode.Value}
}

func filterEventJobDestroyed(resp *etcd.Response) *event.Event {
	if resp.Action != "delete" {
		return nil
	}

	dir, jobName := path.Split(resp.Node.Key)
	dir = strings.TrimSuffix(dir, "/")
	dir, prefixName := path.Split(dir)

	if prefixName != jobPrefix {
		return nil
	}

	return &event.Event{"EventJobDestroyed", jobName, nil}
}

func (r *Registry) jobTargetAgentPath(jobName string) string {
	return path.Join(r.keyPrefix, jobPrefix, jobName, "target")
}

func (r *Registry) jobTargetStatePath(jobName string) string {
	return path.Join(r.keyPrefix, jobPrefix, jobName, "target-state")
}
<commit_msg>fix(registry): lazily migrate legacy payloads<commit_after>package registry

import (
	"errors"
	"path"
	"strings"

	etcdErr "github.com/coreos/fleet/third_party/github.com/coreos/etcd/error"
	"github.com/coreos/fleet/third_party/github.com/coreos/go-etcd/etcd"
	log "github.com/coreos/fleet/third_party/github.com/golang/glog"

	"github.com/coreos/fleet/event"
	"github.com/coreos/fleet/job"
	"github.com/coreos/fleet/unit"
)

const (
	jobPrefix = "job"
)

// GetAllJobs lists all Jobs known by the Registry
func (r *Registry) GetAllJobs() []job.Job {
	var jobs []job.Job

	key := path.Join(r.keyPrefix, jobPrefix)
	resp, err := r.etcd.Get(key, true, true)

	if err != nil {
		return jobs
	}

	for _, dir := range resp.Node.Nodes {
		for _, node := range dir.Nodes {
			if !strings.HasSuffix(node.Key, "object") {
				continue
			}

			j := r.getJobFromJSON(node.Value)
			if j == nil {
				continue
			}

			jobs = append(jobs, *j)
		}
	}

	return jobs
}

// GetJobTarget looks up where the given job is scheduled. If the job has
// been scheduled, the ID the target machine is returned. Otherwise, an
// empty string is returned.
func (r *Registry) GetJobTarget(jobName string) string {
	// Figure out to which Machine this Job is scheduled
	key := r.jobTargetAgentPath(jobName)
	resp, err := r.etcd.Get(key, false, true)
	if err != nil {
		return ""
	}

	return resp.Node.Value
}

func (r *Registry) ClearJobTarget(jobName, machID string) error {
	key := r.jobTargetAgentPath(jobName)
	_, err := r.etcd.CompareAndDelete(key, machID, 0)
	return err
}

// GetJob looks for a Job of the given name in the Registry. It returns a fully
// hydrated Job on success, or nil on any kind of failure.
func (r *Registry) GetJob(jobName string) (j *job.Job) {
	key := path.Join(r.keyPrefix, jobPrefix, jobName, "object")
	resp, err := r.etcd.Get(key, false, true)

	// Assume the error was KeyNotFound and return an empty data structure
	if err != nil {
		return nil
	}

	j = r.getJobFromJSON(resp.Node.Value)
	if j == nil {
		log.V(1).Infof("Error unmarshaling Job(%s): %v", jobName, err)
	}
	return
}

func (r *Registry) getJobFromJSON(val string) *job.Job {
	var jm jobModel
	if err := unmarshal(val, &jm); err != nil {
		return nil
	}

	return r.getJobFromModel(jm)
}

func (r *Registry) getJobFromModel(jm jobModel) *job.Job {
	var err error
	var unit *unit.Unit

	// New-style Jobs should have a populated UnitHash, and the contents of the Unit are stored separately in the Registry
	if !jm.UnitHash.Empty() {
		unit = r.getUnitByHash(jm.UnitHash)
		if unit == nil {
			log.Warningf("No Unit found in Registry for Job(%s)", jm.Name)
			return nil
		}
		if unit.Hash() != jm.UnitHash {
			log.Errorf("Unit Hash %s does not match expected %s for Job(%s)!", unit.Hash(), jm.UnitHash, jm.Name)
			return nil
		}
		log.V(2).Infof("Got Unit for Job(%s) from registry", jm.Name)
	} else {
		// Old-style Jobs had "Payloads" instead of Units, also stored separately in the Registry
		log.V(2).Infof("Legacy Job(%s) has no PayloadHash - looking for associated Payload", jm.Name)
		unit, err = r.getUnitFromLegacyPayload(jm.Name)
		if err != nil {
			log.Errorf("Error retrieving legacy payload for Job(%s)", jm.Name)
			return nil
		} else if unit == nil {
			log.Warningf("No Payload found in Registry for Job(%s)", jm.Name)
			return nil
		}

		log.Infof("Migrating legacy Payload(%s)", jm.Name)
		if err := r.storeOrGetUnit(*unit); err != nil {
			log.Warningf("Unable to migrate legacy Payload: %v", err)
		}
	}

	j := job.NewJob(jm.Name, *unit)

	j.UnitState = r.getUnitState(jm.Name)
	j.State = r.determineJobState(jm.Name)

	return j
}

// jobModel is used for serializing and deserializing Jobs stored in the Registry
type jobModel struct {
	Name     string
	UnitHash unit.Hash
}

// DestroyJob removes a Job object from the repository, along with any legacy
// associated Payload and SignatureSet. It does not yet remove underlying
// Units from the repository.
func (r *Registry) DestroyJob(jobName string) {
	key := path.Join(r.keyPrefix, jobPrefix, jobName)
	r.etcd.Delete(key, true)
	// TODO(jonboulle): add unit reference counting and actually destroying Units
	r.destroyLegacyPayload(jobName)
	r.destroySignatureSetOfJob(jobName)
}

// destroyLegacyPayload removes an old-style Payload from the registry
func (r *Registry) destroyLegacyPayload(payloadName string) {
	key := path.Join(r.keyPrefix, payloadPrefix, payloadName)
	r.etcd.Delete(key, false)
}

// CreateJob attempts to store a Job and its associated Unit in the registry
func (r *Registry) CreateJob(j *job.Job) (err error) {
	if err := r.storeOrGetUnit(j.Unit); err != nil {
		return err
	}

	key := path.Join(r.keyPrefix, jobPrefix, j.Name, "object")

	jm := jobModel{
		Name:     j.Name,
		UnitHash: j.UnitHash,
	}
	json, err := marshal(jm)
	if err != nil {
		return
	}

	_, err = r.etcd.Create(key, json, 0)
	if err != nil && err.(*etcd.EtcdError).ErrorCode == etcdErr.EcodeNodeExist {
		err = errors.New("job already exists")
	}

	return
}

func (r *Registry) GetJobTargetState(jobName string) *job.JobState {
	key := r.jobTargetStatePath(jobName)
	resp, err := r.etcd.Get(key, false, false)
	if err != nil {
		if err.(*etcd.EtcdError).ErrorCode != etcdErr.EcodeNodeExist {
			log.Errorf("Unable to determine target-state of Job(%s): %v", jobName, err)
		}
		return nil
	}

	return job.ParseJobState(resp.Node.Value)
}

func (r *Registry) SetJobTargetState(jobName string, state job.JobState) error {
	key := r.jobTargetStatePath(jobName)
	_, err := r.etcd.Set(key, string(state), 0)
	return err
}

func (es *EventStream) filterJobTargetStateChanges(resp *etcd.Response) *event.Event {
	if resp.Action != "set" {
		return nil
	}

	dir, baseName := path.Split(resp.Node.Key)
	if baseName != "target-state" {
		return nil
	}

	dir = strings.TrimSuffix(dir, "/")
	jobName := path.Base(dir)

	ts := job.ParseJobState(resp.Node.Value)
	if ts == nil {
		return nil
	}

	cs := es.registry.determineJobState(jobName)
	if *cs == *ts {
		return nil
	}

	var cType string
	switch *cs {
	case job.JobStateInactive:
		cType = "CommandLoadJob"
	case job.JobStateLoaded:
		if *ts == job.JobStateInactive {
			cType = "CommandUnloadJob"
		} else if *ts == job.JobStateLaunched {
			cType = "CommandStartJob"
		}
	case job.JobStateLaunched:
		if *ts == job.JobStateLoaded {
			cType = "CommandStopJob"
		} else if *ts == job.JobStateInactive {
			cType = "CommandUnloadJob"
		}
	}

	if cType == "" {
		return nil
	}

	agent := es.registry.GetJobTarget(jobName)
	return &event.Event{cType, jobName, agent}
}

func (r *Registry) ScheduleJob(jobName string, machID string) error {
	key := r.jobTargetAgentPath(jobName)
	_, err := r.etcd.Create(key, machID, 0)
	return err
}

func (r *Registry) LockJob(jobName, context string) *TimedResourceMutex {
	return r.lockResource("job", jobName, context)
}

func filterEventJobScheduled(resp *etcd.Response) *event.Event {
	if resp.Action != "create" {
		return nil
	}

	dir, baseName := path.Split(resp.Node.Key)
	if baseName != "target" {
		return nil
	}

	dir = strings.TrimSuffix(dir, "/")
	dir, jobName := path.Split(dir)

	dir = strings.TrimSuffix(dir, "/")
	dir, prefixName := path.Split(dir)

	if prefixName != jobPrefix {
		return nil
	}

	return &event.Event{"EventJobScheduled", jobName, resp.Node.Value}
}

func filterEventJobUnscheduled(resp *etcd.Response) *event.Event {
	if resp.Action != "delete" && resp.Action != "compareAndDelete" {
		return nil
	}

	dir, baseName := path.Split(resp.Node.Key)
	if baseName != "target" {
		return nil
	}

	dir = strings.TrimSuffix(dir, "/")
	dir, jobName := path.Split(dir)

	dir = strings.TrimSuffix(dir, "/")
	dir, prefixName := path.Split(dir)

	if prefixName != jobPrefix {
		return nil
	}

	if resp.PrevNode == nil {
		return nil
	}

	return &event.Event{"EventJobUnscheduled", jobName, resp.PrevNode.Value}
}

func filterEventJobDestroyed(resp *etcd.Response) *event.Event {
	if resp.Action != "delete" {
		return nil
	}

	dir, jobName := path.Split(resp.Node.Key)
	dir = strings.TrimSuffix(dir, "/")
	dir, prefixName := path.Split(dir)

	if prefixName != jobPrefix {
		return nil
	}

	return &event.Event{"EventJobDestroyed", jobName, nil}
}

func (r *Registry) jobTargetAgentPath(jobName string) string {
	return path.Join(r.keyPrefix, jobPrefix, jobName, "target")
}

func (r *Registry) jobTargetStatePath(jobName string) string {
	return path.Join(r.keyPrefix, jobPrefix, jobName, "target-state")
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
)

const ()

func parse(r io.Reader) (node, error) {
	p := &parser{
		root:   newRootNode(),
		input:  lex(r),
		backup: make([]token, 0, 8),
	}
	if err := p.parse(); err != nil {
		return nil, err
	}
	return p.root, nil
}

type parser struct {
	root   node
	input  chan token
	backup []token
}

func (p *parser) parse() error {
	if p.root == nil {
		p.root = newRootNode()
	}
	return p.root.parse(p)
}

// returns the next token and advances the input stream
func (p *parser) next() token {
	if len(p.backup) > 0 {
		t := p.backup[len(p.backup)-1]
		p.backup = p.backup[:len(p.backup)-1]
		return t
	}
	return <-p.input
}

func (p *parser) peek() token {
	t := p.next()
	p.unread(t)
	return t
}

func (p *parser) unread(t token) {
	if p.backup == nil {
		p.backup = make([]token, 0, 8)
	}
	p.backup = append(p.backup, t)
}

// parse the next value.  This is to be executed in a context where we know we
// want something that is a value to come next, such as after an equals sign.
func (p *parser) parseValue() (interface{}, error) {
	for {
		t := p.next()
		switch t.t {
		case t_error:
			return nil, fmt.Errorf("parse error: saw lex error when looking for value: %v", t.s)
		case t_eof:
			return nil, fmt.Errorf("parse error: unexpected eof when looking for value")
		case t_string:
			return t.s, nil
		case t_list_start:
			l := new(list)

		SIN:
			if p.peek().t == t_list_end {
				p.next()
				return l, nil
			}

			if v, err := p.parseValue(); err != nil {
				return nil, err
			} else {
				l.append(v)
			}

			switch t := p.next(); t.t {
			case t_list_separator:
				goto SIN
			case t_list_end:
				return l, nil
			default:
				return nil, fmt.Errorf("parse error: unexpected %v token while scanning for list", t.t)
			}
		default:
			return nil, fmt.Errorf("parse error: unexpected %v token while looking for value", t.t)
		}
	}
}
<commit_msg>clean<commit_after>package main

import (
	"fmt"
	"io"
)

const ()

func parse(r io.Reader) (node, error) {
	p := &parser{
		root:   newRootNode(),
		input:  lex(r),
		backup: make([]token, 0, 8),
	}
	if err := p.parse(); err != nil {
		return nil, err
	}
	return p.root, nil
}

type parser struct {
	root   node
	input  chan token
	backup []token
}

func (p *parser) parse() error {
	if p.root == nil {
		p.root = newRootNode()
	}
	return p.root.parse(p)
}

// returns the next token and advances the input stream
func (p *parser) next() token {
	if len(p.backup) > 0 {
		t := p.backup[len(p.backup)-1]
		p.backup = p.backup[:len(p.backup)-1]
		return t
	}
	return <-p.input
}

func (p *parser) peek() token {
	t := p.next()
	p.unread(t)
	return t
}

func (p *parser) unread(t token) {
	if p.backup == nil {
		p.backup = make([]token, 0, 8)
	}
	p.backup = append(p.backup, t)
}

// parse the next value.  This is to be executed in a context where we know we
// want something that is a value to come next, such as after an equals sign.
func (p *parser) parseValue() (interface{}, error) {
	for {
		t := p.next()
		switch t.t {
		case t_error:
			return nil, fmt.Errorf("parse error: saw lex error when looking for value: %v", t.s)
		case t_eof:
			return nil, fmt.Errorf("parse error: unexpected eof when looking for value")
		case t_string:
			return t.s, nil
		case t_list_start:
			return p.parseList(new(list))
		default:
			return nil, fmt.Errorf("parse error: unexpected %v token while looking for value", t.t)
		}
	}
}

func (p *parser) parseList(l *list) (*list, error) {
	if p.peek().t == t_list_end {
		p.next()
		return l, nil
	}

	if v, err := p.parseValue(); err != nil {
		return nil, err
	} else {
		l.append(v)
	}

	switch t := p.next(); t.t {
	case t_list_separator:
		return p.parseList(l)
	case t_list_end:
		return l, nil
	default:
		return nil, fmt.Errorf("parse error: unexpected %v token while scanning for list", t.t)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2016, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package sh

import (
	"bufio"
	"fmt"
	"io"
	"regexp"
	"strconv"
)

const (
	_ = -iota
	EOF
	IDENT
	STRING
)

func parse(r io.Reader, name string) error {
	p := &parser{
		r:    bufio.NewReader(r),
		name: name,
		line: 1,
		col:  0,
	}
	p.program()
	return p.err
}

type parser struct {
	r    *bufio.Reader
	tok  int32
	err  error
	name string
	line int
	col  int
}

var reserved = map[rune]bool{
	'\n': true,
	'#':  true,
	'=':  true,
	'&':  true,
	'>':  true,
	'<':  true,
	'|':  true,
	';':  true,
	'(':  true,
	')':  true,
	'{':  true,
	'}':  true,
	'"':  true,
	'\'': true,
}

var space = map[rune]bool{
	' ':  true,
	'\t': true,
}

var ident = regexp.MustCompile(`^[a-zA-Z_]+[a-zA-Z0-9_]*$`)

func (p *parser) next() {
	if p.err != nil {
		return
	}
	r := ' '
	var err error
	for space[r] {
		r, _, err = p.r.ReadRune()
		if err == io.EOF {
			p.tok = EOF
			return
		}
		if err != nil {
			p.errPass(err)
			return
		}
		p.col++
	}
	if reserved[r] {
		if r == '\n' {
			p.line++
			p.col = 0
		}
		p.tok = r
		return
	}
	runes := []rune{r}
	for !reserved[r] && !space[r] {
		r, _, err = p.r.ReadRune()
		if err == io.EOF {
			break
		}
		if err != nil {
			p.errPass(err)
			return
		}
		p.col++
		runes = append(runes, r)
	}
	if len(runes) > 1 {
		if err := p.r.UnreadRune(); err != nil {
			p.errPass(err)
			return
		}
		p.col--
	}
	p.tok = STRING
	s := string(runes)
	if ident.MatchString(s) {
		p.tok = IDENT
	}
	return
}

func (p *parser) discardUpTo(delim byte) {
	_, err := p.r.ReadBytes(delim)
	if err == io.EOF {
		p.tok = EOF
	} else if err != nil {
		p.errPass(err)
	}
}

func (p *parser) got(tok int32) bool {
	if p.tok == tok {
		p.next()
		return true
	}
	return false
}

func (p *parser) want(tok int32) {
	if p.tok != tok {
		p.errWanted(tok)
		return
	}
	p.next()
}

func tokStr(tok int32) string {
	switch tok {
	case EOF:
		return "EOF"
	case STRING:
		return "string"
	case IDENT:
		return "ident"
	default:
		return strconv.QuoteRune(tok)
	}
}

func (p *parser) errPass(err error) {
	p.err = err
	p.tok = EOF
}

func (p *parser) lineErr(format string, v ...interface{}) {
	pos := fmt.Sprintf("%s:%d:%d: ", p.name, p.line, p.col)
	p.errPass(fmt.Errorf(pos+format, v...))
}

func (p *parser) errUnexpected() {
	p.lineErr("unexpected token %s", tokStr(p.tok))
}

func (p *parser) errWanted(tok int32) {
	p.lineErr("unexpected token %s, wanted %s", tokStr(p.tok), tokStr(tok))
}

func (p *parser) program() {
	p.next()
	for p.tok != EOF {
		if p.got('\n') {
			continue
		}
		p.command()
	}
}

func (p *parser) command() {
	switch {
	case p.got('#'):
		p.discardUpTo('\n')
		p.next()
	case p.got('"'):
		p.strContent('"')
	case p.got('\''):
		p.strContent('\'')
	case p.got(STRING):
		fallthrough
	case p.got(IDENT):
		for p.tok != EOF {
			switch {
			case p.got(IDENT):
			case p.got(STRING):
			case p.got('='):
				switch {
				case p.got(IDENT):
				case p.got(STRING):
				}
			case p.got('&'):
				if p.got('&') {
					p.command()
				}
				return
			case p.got('|'):
				p.got('|')
				p.command()
				return
			case p.got('('):
				p.want(')')
				p.want('{')
				for !p.got('}') {
					if p.tok == EOF {
						p.errWanted('}')
						break
					}
					if p.got('\n') {
						continue
					}
					p.command()
				}
				return
			case p.got('>'):
				switch {
				case p.got('>'):
				case p.got('&'):
				}
				p.value()
			case p.got('<'):
				p.value()
			case p.got(';'):
				return
			case p.got('\n'):
				return
			default:
				p.errUnexpected()
			}
		}
	case p.got('{'):
		for !p.got('}') {
			if p.tok == EOF {
				p.errWanted('}')
				break
			}
			if p.got('\n') {
				continue
			}
			p.command()
		}
		switch {
		case p.got(';'):
		case p.got('\n'):
		}
	default:
		p.errUnexpected()
	}
}

func (p *parser) value() {
	switch {
	case p.got(IDENT):
	case p.got(STRING):
	default:
		p.errUnexpected()
	}
}

func (p *parser) strContent(delim byte) {
	_, err := p.r.ReadBytes(delim)
	if err != nil {
		p.errPass(err)
	}
}
<commit_msg>Don't let strings define vars/funcs<commit_after>// Copyright (c) 2016, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package sh

import (
	"bufio"
	"fmt"
	"io"
	"regexp"
	"strconv"
)

const (
	_ = -iota
	EOF
	IDENT
	STRING
)

func parse(r io.Reader, name string) error {
	p := &parser{
		r:    bufio.NewReader(r),
		name: name,
		line: 1,
		col:  0,
	}
	p.program()
	return p.err
}

type parser struct {
	r    *bufio.Reader
	tok  int32
	err  error
	name string
	line int
	col  int
}

var reserved = map[rune]bool{
	'\n': true,
	'#':  true,
	'=':  true,
	'&':  true,
	'>':  true,
	'<':  true,
	'|':  true,
	';':  true,
	'(':  true,
	')':  true,
	'{':  true,
	'}':  true,
	'"':  true,
	'\'': true,
}

var space = map[rune]bool{
	' ':  true,
	'\t': true,
}

var ident = regexp.MustCompile(`^[a-zA-Z_]+[a-zA-Z0-9_]*$`)

func (p *parser) next() {
	if p.err != nil {
		return
	}
	r := ' '
	var err error
	for space[r] {
		r, _, err = p.r.ReadRune()
		if err == io.EOF {
			p.tok = EOF
			return
		}
		if err != nil {
			p.errPass(err)
			return
		}
		p.col++
	}
	if reserved[r] {
		if r == '\n' {
			p.line++
			p.col = 0
		}
		p.tok = r
		return
	}
	runes := []rune{r}
	for !reserved[r] && !space[r] {
		r, _, err = p.r.ReadRune()
		if err == io.EOF {
			break
		}
		if err != nil {
			p.errPass(err)
			return
		}
		p.col++
		runes = append(runes, r)
	}
	if len(runes) > 1 {
		if err := p.r.UnreadRune(); err != nil {
			p.errPass(err)
			return
		}
		p.col--
		runes = runes[:len(runes)-1]
	}
	p.tok = STRING
	s := string(runes)
	if ident.MatchString(s) {
		p.tok = IDENT
	}
	return
}

func (p *parser) discardUpTo(delim byte) {
	_, err := p.r.ReadBytes(delim)
	if err == io.EOF {
		p.tok = EOF
	} else if err != nil {
		p.errPass(err)
	}
}

func (p *parser) got(tok int32) bool {
	if p.tok == tok {
		p.next()
		return true
	}
	return false
}

func (p *parser) want(tok int32) {
	if p.tok != tok {
		p.errWanted(tok)
		return
	}
	p.next()
}

func tokStr(tok int32) string {
	switch tok {
	case EOF:
		return "EOF"
	case STRING:
		return "string"
	case IDENT:
		return "ident"
	default:
		return strconv.QuoteRune(tok)
	}
}

func (p *parser) errPass(err error) {
	p.err = err
	p.tok = EOF
}

func (p *parser) lineErr(format string, v ...interface{}) {
	pos := fmt.Sprintf("%s:%d:%d: ", p.name, p.line, p.col)
	p.errPass(fmt.Errorf(pos+format, v...))
}

func (p *parser) errUnexpected() {
	p.lineErr("unexpected token %s", tokStr(p.tok))
}

func (p *parser) errWanted(tok int32) {
	p.lineErr("unexpected token %s, wanted %s", tokStr(p.tok), tokStr(tok))
}

func (p *parser) program() {
	p.next()
	for p.tok != EOF {
		if p.got('\n') {
			continue
		}
		p.command()
	}
}

func (p *parser) command() {
	switch {
	case p.got('#'):
		p.discardUpTo('\n')
		p.next()
	case p.got('"'):
		p.strContent('"')
	case p.got('\''):
		p.strContent('\'')
	case p.got(IDENT):
		for p.tok != EOF {
			switch {
			case p.got(IDENT):
			case p.got(STRING):
			case p.got('='):
				switch {
				case p.got(IDENT):
				case p.got(STRING):
				}
			case p.got('&'):
				if p.got('&') {
					p.command()
				}
				return
			case p.got('|'):
				p.got('|')
				p.command()
				return
			case p.got('('):
				p.want(')')
				p.want('{')
				for !p.got('}') {
					if p.tok == EOF {
						p.errWanted('}')
						break
					}
					if p.got('\n') {
						continue
					}
					p.command()
				}
				return
			case p.got('>'):
				switch {
				case p.got('>'):
				case p.got('&'):
				}
				p.value()
			case p.got('<'):
				p.value()
			case p.got(';'):
				return
			case p.got('\n'):
				return
			default:
				p.errUnexpected()
			}
		}
	case p.got(STRING):
		for p.tok != EOF {
			switch {
			case p.got(IDENT):
			case p.got(STRING):
			case p.got('&'):
				if p.got('&') {
					p.command()
				}
				return
			case p.got('|'):
				p.got('|')
				p.command()
				return
			case p.got('>'):
				switch {
				case p.got('>'):
				case p.got('&'):
				}
				p.value()
			case p.got('<'):
				p.value()
			case p.got(';'):
				return
			case p.got('\n'):
				return
			default:
				p.errUnexpected()
			}
		}
	case p.got('{'):
		for !p.got('}') {
			if p.tok == EOF {
				p.errWanted('}')
				break
			}
			if p.got('\n') {
				continue
			}
			p.command()
		}
		switch {
		case p.got(';'):
		case p.got('\n'):
		}
	default:
		p.errUnexpected()
	}
}

func (p *parser) value() {
	switch {
	case p.got(IDENT):
	case p.got(STRING):
	default:
		p.errUnexpected()
	}
}

func (p *parser) strContent(delim byte) {
	_, err := p.r.ReadBytes(delim)
	if err != nil {
		p.errPass(err)
	}
}
<|endoftext|>


<commit_before>package arangolite

import "encoding/json"

// DATABASE

// CreateDatabase creates a new database.
type CreateDatabase struct {
	Username string                   `json:"username,omitempty"`
	Name     string                   `json:"name"`
	Extra    json.RawMessage          `json:"extra,omitempty"`
	Passwd   string                   `json:"passwd,omitempty"`
	Active   *bool                    `json:"active,omitempty"`
	Users    []map[string]interface{} `json:"users,omitempty"`
}

func (r *CreateDatabase) Description() string {
	return "CREATE DATABASE"
}

func (r *CreateDatabase) Path() string {
	return "/_api/database"
}

func (r *CreateDatabase) Method() string {
	return "POST"
}

func (r *CreateDatabase) Generate() []byte {
	m, _ := json.Marshal(r)
	return m
}

// DropDatabase deletes a database.
type DropDatabase struct {
	Name string
}

func (r *DropDatabase) Description() string {
	return "DROP DATABASE"
}

func (r *DropDatabase) Path() string {
	return "/_api/database/" + r.Name
}

func (r *DropDatabase) Method() string {
	return "DELETE"
}

func (r *DropDatabase) Generate() []byte {
	return nil
}

// COLLECTION

// CreateCollection creates a collection in database.
type CreateCollection struct {
	JournalSize    int                    `json:"journalSize,omitempty"`
	KeyOptions     map[string]interface{} `json:"keyOptions,omitempty"`
	Name           string                 `json:"name"`
	WaitForSync    *bool                  `json:"waitForSync,omitempty"`
	DoCompact      *bool                  `json:"doCompact,omitempty"`
	IsVolatile     *bool                  `json:"isVolatile,omitempty"`
	ShardKeys      []string               `json:"shardKeys,omitempty"`
	NumberOfShards int                    `json:"numberOfShards,omitempty"`
	IsSystem       *bool                  `json:"isSystem,omitempty"`
	Type           int                    `json:"type,omitempty"`
	IndexBuckets   int                    `json:"indexBuckets,omitempty"`
}

func (r *CreateCollection) Description() string {
	return "CREATE COLLECTION"
}

func (r *CreateCollection) Path() string {
	return "/_api/collection"
}

func (r *CreateCollection) Method() string {
	return "POST"
}

func (r *CreateCollection) Generate() []byte {
	m, _ := json.Marshal(r)
	return m
}

// DropCollection deletes a collection in database.
type DropCollection struct {
	Name string
}

func (r *DropCollection) Description() string {
	return "DROP COLLECTION"
}

func (r *DropCollection) Path() string {
	return "/_api/collection/" + r.Name
}

func (r *DropCollection) Method() string {
	return "DELETE"
}

func (r *DropCollection) Generate() []byte {
	return nil
}

// TruncateCollection deletes a collection in database.
type TruncateCollection struct {
	Name string
}

func (r *TruncateCollection) Description() string {
	return "TRUNCATE COLLECTION"
}

func (r *TruncateCollection) Path() string {
	return "/_api/collection/" + r.Name + "/truncate"
}

func (r *TruncateCollection) Method() string {
	return "PUT"
}

func (r *TruncateCollection) Generate() []byte {
	return nil
}

// INDEX

// CreateHashIndex creates a hash index in database.
type CreateHashIndex struct {
	CollectionName string   `json:"-"`
	Fields         []string `json:"fields,omitempty"`
	Unique         *bool    `json:"unique,omitempty"`
	Type           string   `json:"type,omitempty"`
	Sparse         *bool    `json:"sparse,omitempty"`
}

func (r *CreateHashIndex) Description() string {
	return "CREATE HASH INDEX"
}

func (r *CreateHashIndex) Path() string {
	return "/_api/index?collection=" + r.CollectionName
}

func (r *CreateHashIndex) Method() string {
	return "POST"
}

func (r *CreateHashIndex) Generate() []byte {
	r.Type = "hash"
	m, _ := json.Marshal(r)
	return m
}

// CACHE

// SetCacheProperties sets the query cache properties.
type SetCacheProperties struct {
	Mode       string `json:"mode,omitempty"`
	MaxResults int    `json:"maxResults,omitempty"`
}

func (r *SetCacheProperties) Description() string {
	return "SET CACHE PROPERTIES"
}

func (r *SetCacheProperties) Path() string {
	return "/_api/query-cache/properties"
}

func (r *SetCacheProperties) Method() string {
	return "PUT"
}

func (r *SetCacheProperties) Generate() []byte {
	m, _ := json.Marshal(r)
	return m
}

// GetCacheProperties retrieves the current query cache properties.
type GetCacheProperties struct{}

func (r *GetCacheProperties) Description() string {
	return "GET CACHE PROPERTIES"
}

func (r *GetCacheProperties) Path() string {
	return "/_api/query-cache/properties"
}

func (r *GetCacheProperties) Method() string {
	return "GET"
}

func (r *GetCacheProperties) Generate() []byte {
	return nil
}
<commit_msg>Add new requests to operate on Collections<commit_after>package arangolite

import (
	"encoding/json"
	"strconv"
)

// DATABASE

// CreateDatabase creates a new database.
type CreateDatabase struct {
	Username string                   `json:"username,omitempty"`
	Name     string                   `json:"name"`
	Extra    json.RawMessage          `json:"extra,omitempty"`
	Passwd   string                   `json:"passwd,omitempty"`
	Active   *bool                    `json:"active,omitempty"`
	Users    []map[string]interface{} `json:"users,omitempty"`
}

func (r *CreateDatabase) Description() string {
	return "CREATE DATABASE"
}

func (r *CreateDatabase) Path() string {
	return "/_api/database"
}

func (r *CreateDatabase) Method() string {
	return "POST"
}

func (r *CreateDatabase) Generate() []byte {
	m, _ := json.Marshal(r)
	return m
}

// DropDatabase deletes a database.
type DropDatabase struct {
	Name string
}

func (r *DropDatabase) Description() string {
	return "DROP DATABASE"
}

func (r *DropDatabase) Path() string {
	return "/_api/database/" + r.Name
}

func (r *DropDatabase) Method() string {
	return "DELETE"
}

func (r *DropDatabase) Generate() []byte {
	return nil
}

// COLLECTION

// CreateCollection creates a collection in database.
type CreateCollection struct {
	JournalSize    int                    `json:"journalSize,omitempty"`
	KeyOptions     map[string]interface{} `json:"keyOptions,omitempty"`
	Name           string                 `json:"name"`
	WaitForSync    *bool                  `json:"waitForSync,omitempty"`
	DoCompact      *bool                  `json:"doCompact,omitempty"`
	IsVolatile     *bool                  `json:"isVolatile,omitempty"`
	ShardKeys      []string               `json:"shardKeys,omitempty"`
	NumberOfShards int                    `json:"numberOfShards,omitempty"`
	IsSystem       *bool                  `json:"isSystem,omitempty"`
	Type           int                    `json:"type,omitempty"`
	IndexBuckets   int                    `json:"indexBuckets,omitempty"`
}

func (r *CreateCollection) Description() string {
	return "CREATE COLLECTION"
}

func (r *CreateCollection) Path() string {
	return "/_api/collection"
}

func (r *CreateCollection) Method() string {
	return "POST"
}

func (r *CreateCollection) Generate() []byte {
	m, _ := json.Marshal(r)
	return m
}

// DropCollection deletes a collection in database.
type DropCollection struct {
	Name string
}

func (r *DropCollection) Description() string {
	return "DROP COLLECTION"
}

func (r *DropCollection) Path() string {
	return "/_api/collection/" + r.Name
}

func (r *DropCollection) Method() string {
	return "DELETE"
}

func (r *DropCollection) Generate() []byte {
	return nil
}

// TruncateCollection deletes a collection in database.
type TruncateCollection struct {
	Name string
}

func (r *TruncateCollection) Description() string {
	return "TRUNCATE COLLECTION"
}

func (r *TruncateCollection) Path() string {
	return "/_api/collection/" + r.Name + "/truncate"
}

func (r *TruncateCollection) Method() string {
	return "PUT"
}

func (r *TruncateCollection) Generate() []byte {
	return nil
}

// INDEX

// CreateHashIndex creates a hash index in database.
type CreateHashIndex struct {
	CollectionName string   `json:"-"`
	Fields         []string `json:"fields,omitempty"`
	Unique         *bool    `json:"unique,omitempty"`
	Type           string   `json:"type,omitempty"`
	Sparse         *bool    `json:"sparse,omitempty"`
}

func (r *CreateHashIndex) Description() string {
	return "CREATE HASH INDEX"
}

func (r *CreateHashIndex) Path() string {
	return "/_api/index?collection=" + r.CollectionName
}

func (r *CreateHashIndex) Method() string {
	return "POST"
}

func (r *CreateHashIndex) Generate() []byte {
	r.Type = "hash"
	m, _ := json.Marshal(r)
	return m
}

// CACHE

// SetCacheProperties sets the query cache properties.
type SetCacheProperties struct {
	Mode       string `json:"mode,omitempty"`
	MaxResults int    `json:"maxResults,omitempty"`
}

func (r *SetCacheProperties) Description() string {
	return "SET CACHE PROPERTIES"
}

func (r *SetCacheProperties) Path() string {
	return "/_api/query-cache/properties"
}

func (r *SetCacheProperties) Method() string {
	return "PUT"
}

func (r *SetCacheProperties) Generate() []byte {
	m, _ := json.Marshal(r)
	return m
}

// GetCacheProperties retrieves the current query cache properties.
type GetCacheProperties struct{}

func (r *GetCacheProperties) Description() string {
	return "GET CACHE PROPERTIES"
}

func (r *GetCacheProperties) Path() string {
	return "/_api/query-cache/properties"
}

func (r *GetCacheProperties) Method() string {
	return "GET"
}

func (r *GetCacheProperties) Generate() []byte {
	return nil
}

type CollectionInfo struct {
	Id string        `json:"id"`
	Name string       `json:"name"`
	IsSystem bool       `json:"isSystem"`
	Status int        `json:"status"`
	Type int          `json:"type"`
}

type CollectionInfoList struct {
	Collections []CollectionInfo    `json:"collections"`
	Error bool                      `json:"error"`
	Code int      `json:"code"`
}

// ListCollections lists all collections from the current DB
type ListCollections struct {
	includeSystem bool
}

func (c *ListCollections) Description() string {
	return "LIST COLLECTIONS"
}

func (c *ListCollections) Path() string {
	return "/_api/collection?excludeSystem=" + strconv.FormatBool(!c.includeSystem)
}

func (c *ListCollections) Method() string {
	return "GET"
}

func (c *ListCollections) Generate() []byte {
	return nil
}

// CollectionInfo gets information about the collection
type GetCollectionInfo struct {
	CollectionName string
	IncludeSystem bool
}

func (c *GetCollectionInfo) Description() string {
	return "COLLECTION INFO"
}

func (c *GetCollectionInfo) Path() string {
	return "/_api/collection/" + c.CollectionName + "?excludeSystem=" + strconv.FormatBool(!c.IncludeSystem)
}

func (c *GetCollectionInfo) Method() string {
	return "GET"
}

func (c *GetCollectionInfo) Generate() []byte {
	return nil
}

// ImportCollection imports data to a collection
type ImportCollection struct {
	CollectionName string
	Data           []byte
}

func (c *ImportCollection) Description() string {
	return "IMPORT COLLECTION"
}

func (c *ImportCollection) Path() string {
	return "/_api/import/?type=auto&collection=" + c.CollectionName
}

func (c *ImportCollection) Method() string {
	return "POST"
}

func (c *ImportCollection) Generate() []byte {
	return c.Data
}<|endoftext|>