<commit_before>package testutil

// TestServer is a test helper. It uses a fork/exec model to create
// a test Nomad server instance in the background and initialize it
// with some data and/or services. The test server can then be used
// to run a unit test, and offers an easy API to tear itself down
// when the test has completed. The only prerequisite is to have a nomad
// binary available on the $PATH.
//
// This package does not use Nomad's official API client. This is
// because we use TestServer to test the API client, which would
// otherwise cause an import cycle.

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"sync/atomic"
	"testing"

	"github.com/hashicorp/go-cleanhttp"
)

// offset is used to atomically increment the port numbers.
var offset uint64

// TestServerConfig is the main server configuration struct.
type TestServerConfig struct {
	NodeName          string        `json:"name,omitempty"`
	DataDir           string        `json:"data_dir,omitempty"`
	Region            string        `json:"region,omitempty"`
	DisableCheckpoint bool          `json:"disable_update_check"`
	LogLevel          string        `json:"log_level,omitempty"`
	Ports             *PortsConfig  `json:"ports,omitempty"`
	Server            *ServerConfig `json:"server,omitempty"`
	Client            *ClientConfig `json:"client,omitempty"`
	DevMode           bool          `json:"-"`
	Stdout, Stderr    io.Writer     `json:"-"`
}

// PortsConfig is used to configure the network ports we use.
type PortsConfig struct {
	HTTP int `json:"http,omitempty"`
	RPC  int `json:"rpc,omitempty"`
	Serf int `json:"serf,omitempty"`
}

// ServerConfig is used to configure the nomad server.
type ServerConfig struct {
	Enabled         bool `json:"enabled"`
	BootstrapExpect int  `json:"bootstrap_expect"`
}

// ClientConfig is used to configure the client
type ClientConfig struct {
	Enabled bool `json:"enabled"`
}

// ServerConfigCallback is a function interface which can be
// passed to NewTestServerConfig to modify the server config.
type ServerConfigCallback func(c *TestServerConfig)

// defaultServerConfig returns a new TestServerConfig struct
// with all of the listen ports incremented by one.
func defaultServerConfig() *TestServerConfig {
	idx := int(atomic.AddUint64(&offset, 1))

	return &TestServerConfig{
		NodeName:          fmt.Sprintf("node%d", idx),
		DisableCheckpoint: true,
		LogLevel:          "DEBUG",
		Ports: &PortsConfig{
			HTTP: 20000 + idx,
			RPC:  21000 + idx,
			Serf: 22000 + idx,
		},
		Server: &ServerConfig{
			Enabled:         true,
			BootstrapExpect: 1,
		},
		Client: &ClientConfig{
			Enabled: false,
		},
	}
}

// TestServer is the main server wrapper struct.
type TestServer struct {
	PID    int
	Config *TestServerConfig
	t      *testing.T

	HTTPAddr   string
	SerfAddr   string
	HTTPClient *http.Client
}

// NewTestServer creates a new TestServer, and makes a call to
// an optional callback function to modify the configuration.
func NewTestServer(t *testing.T, cb ServerConfigCallback) *TestServer {
	if path, err := exec.LookPath("nomad"); err != nil || path == "" {
		t.Skip("nomad not found on $PATH, skipping")
	}

	dataDir, err := ioutil.TempDir("", "nomad")
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	configFile, err := ioutil.TempFile(dataDir, "nomad")
	if err != nil {
		defer os.RemoveAll(dataDir)
		t.Fatalf("err: %s", err)
	}

	nomadConfig := defaultServerConfig()
	nomadConfig.DataDir = dataDir

	if cb != nil {
		cb(nomadConfig)
	}

	configContent, err := json.Marshal(nomadConfig)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	if _, err := configFile.Write(configContent); err != nil {
		t.Fatalf("err: %s", err)
	}
	configFile.Close()

	stdout := io.Writer(os.Stdout)
	if nomadConfig.Stdout != nil {
		stdout = nomadConfig.Stdout
	}

	stderr := io.Writer(os.Stderr)
	if nomadConfig.Stderr != nil {
		stderr = nomadConfig.Stderr
	}

	args := []string{"agent", "-config", configFile.Name()}
	if nomadConfig.DevMode {
		args = append(args, "-dev")
	}

	// Start the server
	cmd := exec.Command("nomad", args...)
	cmd.Stdout = stdout
	cmd.Stderr = stderr
	if err := cmd.Start(); err != nil {
		t.Fatalf("err: %s", err)
	}

	client := cleanhttp.DefaultClient()

	server := &TestServer{
		Config: nomadConfig,
		PID:    cmd.Process.Pid,
		t:      t,

		HTTPAddr:   fmt.Sprintf("127.0.0.1:%d", nomadConfig.Ports.HTTP),
		SerfAddr:   fmt.Sprintf("127.0.0.1:%d", nomadConfig.Ports.Serf),
		HTTPClient: client,
	}

	// Wait for the server to be ready
	if nomadConfig.Server.Enabled && nomadConfig.Server.BootstrapExpect != 0 {
		server.waitForLeader()
	} else {
		server.waitForAPI()
	}
	return server
}

// Stop stops the test Nomad server, and removes the Nomad data
// directory once we are done.
func (s *TestServer) Stop() {
	defer os.RemoveAll(s.Config.DataDir)

	cmd := exec.Command("kill", "-9", fmt.Sprintf("%d", s.PID))
	if err := cmd.Run(); err != nil {
		s.t.Errorf("err: %s", err)
	}
}

// waitForAPI waits for only the agent HTTP endpoint to start
// responding. This is an indication that the agent has started,
// but will likely return before a leader is elected.
func (s *TestServer) waitForAPI() {
	WaitForResult(func() (bool, error) {
		resp, err := s.HTTPClient.Get(s.url("/v1/agent/self"))
		if err != nil {
			return false, err
		}
		defer resp.Body.Close()
		if err := s.requireOK(resp); err != nil {
			return false, err
		}
		return true, nil
	}, func(err error) {
		defer s.Stop()
		s.t.Fatalf("err: %s", err)
	})
}

// waitForLeader waits for the Nomad server's HTTP API to become
// available, and then waits for a known leader and an index of
// 1 or more to be observed to confirm leader election is done.
func (s *TestServer) waitForLeader() {
	WaitForResult(func() (bool, error) {
		// Query the API and check the status code
		resp, err := s.HTTPClient.Get(s.url("/v1/jobs"))
		if err != nil {
			return false, err
		}
		defer resp.Body.Close()
		if err := s.requireOK(resp); err != nil {
			return false, err
		}

		// Ensure we have a leader and a node registeration
		if leader := resp.Header.Get("X-Nomad-KnownLeader"); leader != "true" {
			return false, fmt.Errorf("Nomad leader status: %#v", leader)
		}
		return true, nil
	}, func(err error) {
		defer s.Stop()
		s.t.Fatalf("err: %s", err)
	})
}

// url is a helper function which takes a relative URL and
// makes it into a proper URL against the local Nomad server.
func (s *TestServer) url(path string) string {
	return fmt.Sprintf("http://%s%s", s.HTTPAddr, path)
}

// requireOK checks the HTTP response code and ensures it is acceptable.
func (s *TestServer) requireOK(resp *http.Response) error {
	if resp.StatusCode != 200 {
		return fmt.Errorf("Bad status code: %d", resp.StatusCode)
	}
	return nil
}

// put performs a new HTTP PUT request.
func (s *TestServer) put(path string, body io.Reader) *http.Response {
	req, err := http.NewRequest("PUT", s.url(path), body)
	if err != nil {
		s.t.Fatalf("err: %s", err)
	}
	resp, err := s.HTTPClient.Do(req)
	if err != nil {
		s.t.Fatalf("err: %s", err)
	}
	if err := s.requireOK(resp); err != nil {
		defer resp.Body.Close()
		s.t.Fatal(err)
	}
	return resp
}

// get performs a new HTTP GET request.
func (s *TestServer) get(path string) *http.Response {
	resp, err := s.HTTPClient.Get(s.url(path))
	if err != nil {
		s.t.Fatalf("err: %s", err)
	}
	if err := s.requireOK(resp); err != nil {
		defer resp.Body.Close()
		s.t.Fatal(err)
	}
	return resp
}

// encodePayload returns a new io.Reader wrapping the encoded contents
// of the payload, suitable for passing directly to a new request.
func (s *TestServer) encodePayload(payload interface{}) io.Reader {
	var encoded bytes.Buffer
	enc := json.NewEncoder(&encoded)
	if err := enc.Encode(payload); err != nil {
		s.t.Fatalf("err: %s", err)
	}
	return &encoded
}
<commit_msg>Make testutil.TestServer work correctly on Windows.<commit_after>package testutil

// TestServer is a test helper. It uses a fork/exec model to create
// a test Nomad server instance in the background and initialize it
// with some data and/or services. The test server can then be used
// to run a unit test, and offers an easy API to tear itself down
// when the test has completed. The only prerequisite is to have a nomad
// binary available on the $PATH.
//
// This package does not use Nomad's official API client. This is
// because we use TestServer to test the API client, which would
// otherwise cause an import cycle.

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"sync/atomic"
	"testing"

	"github.com/hashicorp/go-cleanhttp"
)

// offset is used to atomically increment the port numbers.
var offset uint64

// TestServerConfig is the main server configuration struct.
type TestServerConfig struct {
	NodeName          string        `json:"name,omitempty"`
	DataDir           string        `json:"data_dir,omitempty"`
	Region            string        `json:"region,omitempty"`
	DisableCheckpoint bool          `json:"disable_update_check"`
	LogLevel          string        `json:"log_level,omitempty"`
	Ports             *PortsConfig  `json:"ports,omitempty"`
	Server            *ServerConfig `json:"server,omitempty"`
	Client            *ClientConfig `json:"client,omitempty"`
	DevMode           bool          `json:"-"`
	Stdout, Stderr    io.Writer     `json:"-"`
}

// PortsConfig is used to configure the network ports we use.
type PortsConfig struct {
	HTTP int `json:"http,omitempty"`
	RPC  int `json:"rpc,omitempty"`
	Serf int `json:"serf,omitempty"`
}

// ServerConfig is used to configure the nomad server.
type ServerConfig struct {
	Enabled         bool `json:"enabled"`
	BootstrapExpect int  `json:"bootstrap_expect"`
}

// ClientConfig is used to configure the client
type ClientConfig struct {
	Enabled bool `json:"enabled"`
}

// ServerConfigCallback is a function interface which can be
// passed to NewTestServerConfig to modify the server config.
type ServerConfigCallback func(c *TestServerConfig)

// defaultServerConfig returns a new TestServerConfig struct
// with all of the listen ports incremented by one.
func defaultServerConfig() *TestServerConfig {
	idx := int(atomic.AddUint64(&offset, 1))

	return &TestServerConfig{
		NodeName:          fmt.Sprintf("node%d", idx),
		DisableCheckpoint: true,
		LogLevel:          "DEBUG",
		Ports: &PortsConfig{
			HTTP: 20000 + idx,
			RPC:  21000 + idx,
			Serf: 22000 + idx,
		},
		Server: &ServerConfig{
			Enabled:         true,
			BootstrapExpect: 1,
		},
		Client: &ClientConfig{
			Enabled: false,
		},
	}
}

// TestServer is the main server wrapper struct.
type TestServer struct {
	cmd    *exec.Cmd
	Config *TestServerConfig
	t      *testing.T

	HTTPAddr   string
	SerfAddr   string
	HTTPClient *http.Client
}

// NewTestServer creates a new TestServer, and makes a call to
// an optional callback function to modify the configuration.
func NewTestServer(t *testing.T, cb ServerConfigCallback) *TestServer {
	if path, err := exec.LookPath("nomad"); err != nil || path == "" {
		t.Skip("nomad not found on $PATH, skipping")
	}

	dataDir, err := ioutil.TempDir("", "nomad")
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	configFile, err := ioutil.TempFile(dataDir, "nomad")
	if err != nil {
		defer os.RemoveAll(dataDir)
		t.Fatalf("err: %s", err)
	}
	defer configFile.Close()

	nomadConfig := defaultServerConfig()
	nomadConfig.DataDir = dataDir

	if cb != nil {
		cb(nomadConfig)
	}

	configContent, err := json.Marshal(nomadConfig)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	if _, err := configFile.Write(configContent); err != nil {
		t.Fatalf("err: %s", err)
	}
	configFile.Close()

	stdout := io.Writer(os.Stdout)
	if nomadConfig.Stdout != nil {
		stdout = nomadConfig.Stdout
	}

	stderr := io.Writer(os.Stderr)
	if nomadConfig.Stderr != nil {
		stderr = nomadConfig.Stderr
	}

	args := []string{"agent", "-config", configFile.Name()}
	if nomadConfig.DevMode {
		args = append(args, "-dev")
	}

	// Start the server
	cmd := exec.Command("nomad", args...)
	cmd.Stdout = stdout
	cmd.Stderr = stderr
	if err := cmd.Start(); err != nil {
		t.Fatalf("err: %s", err)
	}

	client := cleanhttp.DefaultClient()

	server := &TestServer{
		Config: nomadConfig,
		cmd:    cmd,
		t:      t,

		HTTPAddr:   fmt.Sprintf("127.0.0.1:%d", nomadConfig.Ports.HTTP),
		SerfAddr:   fmt.Sprintf("127.0.0.1:%d", nomadConfig.Ports.Serf),
		HTTPClient: client,
	}

	// Wait for the server to be ready
	if nomadConfig.Server.Enabled && nomadConfig.Server.BootstrapExpect != 0 {
		server.waitForLeader()
	} else {
		server.waitForAPI()
	}
	return server
}

// Stop stops the test Nomad server, and removes the Nomad data
// directory once we are done.
func (s *TestServer) Stop() {
	defer os.RemoveAll(s.Config.DataDir)

	if err := s.cmd.Process.Kill(); err != nil {
		s.t.Errorf("err: %s", err)
	}

	// wait for the process to exit to be sure that the data dir can be
	// deleted on all platforms.
	s.cmd.Wait()
}

// waitForAPI waits for only the agent HTTP endpoint to start
// responding. This is an indication that the agent has started,
// but will likely return before a leader is elected.
func (s *TestServer) waitForAPI() {
	WaitForResult(func() (bool, error) {
		resp, err := s.HTTPClient.Get(s.url("/v1/agent/self"))
		if err != nil {
			return false, err
		}
		defer resp.Body.Close()
		if err := s.requireOK(resp); err != nil {
			return false, err
		}
		return true, nil
	}, func(err error) {
		defer s.Stop()
		s.t.Fatalf("err: %s", err)
	})
}

// waitForLeader waits for the Nomad server's HTTP API to become
// available, and then waits for a known leader and an index of
// 1 or more to be observed to confirm leader election is done.
func (s *TestServer) waitForLeader() {
	WaitForResult(func() (bool, error) {
		// Query the API and check the status code
		resp, err := s.HTTPClient.Get(s.url("/v1/jobs"))
		if err != nil {
			return false, err
		}
		defer resp.Body.Close()
		if err := s.requireOK(resp); err != nil {
			return false, err
		}

		// Ensure we have a leader and a node registeration
		if leader := resp.Header.Get("X-Nomad-KnownLeader"); leader != "true" {
			return false, fmt.Errorf("Nomad leader status: %#v", leader)
		}
		return true, nil
	}, func(err error) {
		defer s.Stop()
		s.t.Fatalf("err: %s", err)
	})
}

// url is a helper function which takes a relative URL and
// makes it into a proper URL against the local Nomad server.
func (s *TestServer) url(path string) string {
	return fmt.Sprintf("http://%s%s", s.HTTPAddr, path)
}

// requireOK checks the HTTP response code and ensures it is acceptable.
func (s *TestServer) requireOK(resp *http.Response) error {
	if resp.StatusCode != 200 {
		return fmt.Errorf("Bad status code: %d", resp.StatusCode)
	}
	return nil
}

// put performs a new HTTP PUT request.
func (s *TestServer) put(path string, body io.Reader) *http.Response {
	req, err := http.NewRequest("PUT", s.url(path), body)
	if err != nil {
		s.t.Fatalf("err: %s", err)
	}
	resp, err := s.HTTPClient.Do(req)
	if err != nil {
		s.t.Fatalf("err: %s", err)
	}
	if err := s.requireOK(resp); err != nil {
		defer resp.Body.Close()
		s.t.Fatal(err)
	}
	return resp
}

// get performs a new HTTP GET request.
func (s *TestServer) get(path string) *http.Response {
	resp, err := s.HTTPClient.Get(s.url(path))
	if err != nil {
		s.t.Fatalf("err: %s", err)
	}
	if err := s.requireOK(resp); err != nil {
		defer resp.Body.Close()
		s.t.Fatal(err)
	}
	return resp
}

// encodePayload returns a new io.Reader wrapping the encoded contents
// of the payload, suitable for passing directly to a new request.
func (s *TestServer) encodePayload(payload interface{}) io.Reader {
	var encoded bytes.Buffer
	enc := json.NewEncoder(&encoded)
	if err := enc.Encode(payload); err != nil {
		s.t.Fatalf("err: %s", err)
	}
	return &encoded
}
<|endoftext|>


<commit_before>package main

import (
  "fmt"
  "math/rand"
  "sync"
  "time"
)

var wg sync.WaitGroup

func main() {
  wg.Add(2)

  fmt.Println("Start Goroutines")s
  go printCounts("A")
  go printCounts("B")
  fmt.Println("Waiting To Finish")
  wg.Wait()
  fmt.Println("\nTerminating Program")
}

func printCounts(label string) {
  defer wg.Done()
  for count := 1; count <= 10; count++ {
    sleep := rand.Int63n(1000)
    time.Sleep(time.Duration(sleep) * time.Millisecond)
    fmt.Println("Count: %d from %s\n", count, label)
  }
}
<commit_msg>update<commit_after>package main

import (
  "fmt"
  "math/rand"
  "sync"
  "time"
)

var wg sync.WaitGroup

func main() {
  wg.Add(2)

  fmt.Println("Start Goroutines")
  go printCounts("A")
  go printCounts("B")
  fmt.Println("Waiting To Finish")
  wg.Wait()
  fmt.Println("\nTerminating Program")
}

func printCounts(label string) {
  defer wg.Done()
  for count := 1; count <= 10; count++ {
    sleep := rand.Int63n(1000)
    time.Sleep(time.Duration(sleep) * time.Millisecond)
    fmt.Println("Count: %d from %s\n", count, label)
  }
}
<|endoftext|>


<commit_before>// Package auth manages auth tokens for command line usage.
package auth

import (
	"bufio"
	"context"
	"golang.org/x/crypto/ssh/terminal"
	"io"
	"log"
	"os"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/ptypes"
	tspb "github.com/golang/protobuf/ptypes/timestamp"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	gstatus "google.golang.org/grpc/status"

	"pixur.org/pixur/api"
	"pixur.org/pixur/be/status"
	"pixur.org/pixur/tools/config"
)

const (
	authPwtHeaderKey = "pixur-auth-token"
	pixPwtHeaderKey  = "pixur-pix-token"
	httpHeaderKey    = "pixur-http-header-bin"
)

// cc is optional
func Auth(ctx context.Context, cc **grpc.ClientConn) (context.Context, status.S) {
	conf, sts := getAndUpdateConfig(ctx, cc)
	if sts != nil {
		return nil, sts
	}

	md := make(metadata.MD, 2)
	if conf.AuthToken != "" {
		md.Set(authPwtHeaderKey, conf.AuthToken)
	}
	if conf.PixToken != "" {
		md.Set(pixPwtHeaderKey, conf.PixToken)
	}
	return metadata.NewOutgoingContext(ctx, md), nil
}

func getAndUpdateConfig(ctx context.Context, cc **grpc.ClientConn) (
	_ *config.Config, stscap status.S) {
	oldconf, sts := config.GetConfig()
	if sts != nil {
		if sts.Code() == codes.NotFound {
			oldconf = &config.Config{}
		} else {
			return nil, sts
		}
	}

	newconf := proto.Clone(oldconf).(*config.Config)
	if sts := updateConfig(ctx, cc, newconf); sts != nil {
		return nil, sts
	}
	if !proto.Equal(oldconf, newconf) {
		if sts := config.SetConfig(newconf); sts != nil {
			status.ReplaceOrSuppress(&stscap, sts)
		}
	}

	return newconf, nil
}

// shouldRefresh determines if a token refresh is needed.  notAfter indicates when the token is
// expired and must be non-nil.  softNotAfter indicates that the token is still valid, but should be
// refreshed.  softNotAfter may be nil.  The return value useToken indicates that the auth token
// can be reusued.  useCreds indicates if a ident/secret is needed.  At most one of useToken or
// useCreds can be true.
func shouldRefresh(now func() time.Time, notAfter, softNotAfter *tspb.Timestamp) (
	useToken, useCreds bool, _ status.S) {
	if notAfter == nil {
		return false, false, status.Internal(nil, "missing not after")
	}
	na, err := ptypes.Timestamp(notAfter)
	if err != nil {
		return false, false, status.Internal(err, "bad not after")
	}
	ts := now()
	if ts.After(na) {
		return false, true, nil
	}
	if softNotAfter != nil {
		sna, err := ptypes.Timestamp(softNotAfter)
		if err != nil {
			return false, false, status.Internal(err, "bad soft not after")
		}
		if ts.After(sna) {
			return true, false, nil
		}
	}
	return false, false, nil
}

func updateConfig(ctx context.Context, cc **grpc.ClientConn, conf *config.Config) (stscap status.S) {
	var (
		useToken bool
		useCreds bool = true
	)
	if conf.AuthPayload != nil || conf.AuthToken == "" {
		var sts status.S
		useToken, useCreds, sts =
			shouldRefresh(time.Now, conf.AuthPayload.NotAfter, conf.AuthPayload.SoftNotAfter)
		if sts != nil {
			return sts
		}
		if !useToken && !useCreds {
			return nil
		}
	}

	var conn *grpc.ClientConn
	if cc == nil || *cc == nil {
		if conf.PixurTarget == "" {
			target, sts := getTargetFromTerm(os.Stderr, os.Stdin)
			if sts != nil {
				return sts
			}
			conf.PixurTarget = target
		}
		var err error
		conn, err = grpc.DialContext(ctx, conf.PixurTarget, grpc.WithInsecure())
		if err != nil {
			return status.From(err)
		}
		if cc != nil {
			*cc = conn
		} else {
			defer func() {

				if err := conn.Close(); err != nil {
					status.ReplaceOrSuppress(&stscap, status.From(err))

				}
			}()
		}

	}

	client := api.NewPixurServiceClient(conn)

	if useToken && conf.AuthToken != "" {
		res, err := client.GetRefreshToken(ctx, &api.GetRefreshTokenRequest{
			PreviousAuthToken: conf.AuthToken,
		})
		if err != nil {
			if gstatus.Code(err) == codes.Unauthenticated {
				log.Println("Previous Auth token denied, falling back to creds", err)
				useCreds = true
			} else {
				return status.From(err)
			}
		} else {
			conf.AuthToken, conf.PixToken = res.AuthToken, res.PixToken
			conf.AuthPayload, conf.PixPayload = res.AuthPayload, res.PixPayload
			return nil
		}
	}

	if useCreds {
		ident, secret, sts := getIdentSecretFromTerm(os.Stderr, os.Stdin)
		if sts != nil {
			return sts
		}
		res, err := client.GetRefreshToken(ctx, &api.GetRefreshTokenRequest{
			Ident:  ident,
			Secret: secret,
		})
		if err != nil {
			return status.From(err)
		} else {
			conf.AuthToken, conf.PixToken = res.AuthToken, res.PixToken
			conf.AuthPayload, conf.PixPayload = res.AuthPayload, res.PixPayload
			return nil
		}
	}

	panic("unreachable")
}

func getTargetFromTerm(w io.Writer, r io.Reader) (string, status.S) {
	if _, err := w.Write([]byte("gRPC Target of Pixur Backend (e.g. dns:///localhost:8079): \n")); err != nil {
		return "", status.From(err)
	}
	br := bufio.NewReader(r)
	target, err := br.ReadString('\n')
	if err != nil {
		return "", status.From(err)
	}

	return target[:len(target)-1], nil
}

func getIdentSecretFromTerm(w io.Writer, r *os.File) (string, string, status.S) {
	if _, err := w.Write([]byte("User Ident (e.g. admin): \n")); err != nil {
		return "", "", status.From(err)
	}
	br := bufio.NewReader(r)
	ident, err := br.ReadString('\n')
	if err != nil {
		return "", "", status.From(err)
	}

	if _, err := w.Write([]byte("Secret (e.g. 12345): \n")); err != nil {
		return "", "", status.From(err)
	}
	secret, err := terminal.ReadPassword(int(r.Fd()))
	if err != nil {
		return "", "", status.From(err)
	}

	return ident[:len(ident)-1], string(secret), nil
}
<commit_msg>tools/auth: remove unneeded condition in token refresh<commit_after>// Package auth manages auth tokens for command line usage.
package auth

import (
	"bufio"
	"context"
	"golang.org/x/crypto/ssh/terminal"
	"io"
	"log"
	"os"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/ptypes"
	tspb "github.com/golang/protobuf/ptypes/timestamp"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	gstatus "google.golang.org/grpc/status"

	"pixur.org/pixur/api"
	"pixur.org/pixur/be/status"
	"pixur.org/pixur/tools/config"
)

const (
	authPwtHeaderKey = "pixur-auth-token"
	pixPwtHeaderKey  = "pixur-pix-token"
	httpHeaderKey    = "pixur-http-header-bin"
)

// cc is optional
func Auth(ctx context.Context, cc **grpc.ClientConn) (context.Context, status.S) {
	conf, sts := getAndUpdateConfig(ctx, cc)
	if sts != nil {
		return nil, sts
	}

	md := make(metadata.MD, 2)
	if conf.AuthToken != "" {
		md.Set(authPwtHeaderKey, conf.AuthToken)
	}
	if conf.PixToken != "" {
		md.Set(pixPwtHeaderKey, conf.PixToken)
	}
	return metadata.NewOutgoingContext(ctx, md), nil
}

func getAndUpdateConfig(ctx context.Context, cc **grpc.ClientConn) (
	_ *config.Config, stscap status.S) {
	oldconf, sts := config.GetConfig()
	if sts != nil {
		if sts.Code() == codes.NotFound {
			oldconf = &config.Config{}
		} else {
			return nil, sts
		}
	}

	newconf := proto.Clone(oldconf).(*config.Config)
	if sts := updateConfig(ctx, cc, newconf); sts != nil {
		return nil, sts
	}
	if !proto.Equal(oldconf, newconf) {
		if sts := config.SetConfig(newconf); sts != nil {
			status.ReplaceOrSuppress(&stscap, sts)
		}
	}

	return newconf, nil
}

// shouldRefresh determines if a token refresh is needed.  notAfter indicates when the token is
// expired and must be non-nil.  softNotAfter indicates that the token is still valid, but should be
// refreshed.  softNotAfter may be nil.  The return value useToken indicates that the auth token
// can be reusued.  useCreds indicates if a ident/secret is needed.  At most one of useToken or
// useCreds can be true.
func shouldRefresh(now func() time.Time, notAfter, softNotAfter *tspb.Timestamp) (
	useToken, useCreds bool, _ status.S) {
	if notAfter == nil {
		return false, false, status.Internal(nil, "missing not after")
	}
	na, err := ptypes.Timestamp(notAfter)
	if err != nil {
		return false, false, status.Internal(err, "bad not after")
	}
	ts := now()
	if ts.After(na) {
		return false, true, nil
	}
	if softNotAfter != nil {
		sna, err := ptypes.Timestamp(softNotAfter)
		if err != nil {
			return false, false, status.Internal(err, "bad soft not after")
		}
		if ts.After(sna) {
			return true, false, nil
		}
	}
	return false, false, nil
}

func updateConfig(ctx context.Context, cc **grpc.ClientConn, conf *config.Config) (stscap status.S) {
	var (
		useToken bool
		useCreds bool = true
	)
	if conf.AuthPayload != nil {
		var sts status.S
		useToken, useCreds, sts =
			shouldRefresh(time.Now, conf.AuthPayload.NotAfter, conf.AuthPayload.SoftNotAfter)
		if sts != nil {
			return sts
		}
		if !useToken && !useCreds {
			return nil
		}
	}

	var conn *grpc.ClientConn
	if cc == nil || *cc == nil {
		if conf.PixurTarget == "" {
			target, sts := getTargetFromTerm(os.Stderr, os.Stdin)
			if sts != nil {
				return sts
			}
			conf.PixurTarget = target
		}
		var err error
		conn, err = grpc.DialContext(ctx, conf.PixurTarget, grpc.WithInsecure())
		if err != nil {
			return status.From(err)
		}
		if cc != nil {
			*cc = conn
		} else {
			defer func() {

				if err := conn.Close(); err != nil {
					status.ReplaceOrSuppress(&stscap, status.From(err))

				}
			}()
		}

	}

	client := api.NewPixurServiceClient(conn)

	if useToken && conf.AuthToken != "" {
		res, err := client.GetRefreshToken(ctx, &api.GetRefreshTokenRequest{
			PreviousAuthToken: conf.AuthToken,
		})
		if err != nil {
			if gstatus.Code(err) == codes.Unauthenticated {
				log.Println("Previous Auth token denied, falling back to creds", err)
				useCreds = true
			} else {
				return status.From(err)
			}
		} else {
			conf.AuthToken, conf.PixToken = res.AuthToken, res.PixToken
			conf.AuthPayload, conf.PixPayload = res.AuthPayload, res.PixPayload
			return nil
		}
	}

	if useCreds {
		ident, secret, sts := getIdentSecretFromTerm(os.Stderr, os.Stdin)
		if sts != nil {
			return sts
		}
		res, err := client.GetRefreshToken(ctx, &api.GetRefreshTokenRequest{
			Ident:  ident,
			Secret: secret,
		})
		if err != nil {
			return status.From(err)
		} else {
			conf.AuthToken, conf.PixToken = res.AuthToken, res.PixToken
			conf.AuthPayload, conf.PixPayload = res.AuthPayload, res.PixPayload
			return nil
		}
	}

	panic("unreachable")
}

func getTargetFromTerm(w io.Writer, r io.Reader) (string, status.S) {
	if _, err := w.Write([]byte("gRPC Target of Pixur Backend (e.g. dns:///localhost:8079): \n")); err != nil {
		return "", status.From(err)
	}
	br := bufio.NewReader(r)
	target, err := br.ReadString('\n')
	if err != nil {
		return "", status.From(err)
	}

	return target[:len(target)-1], nil
}

func getIdentSecretFromTerm(w io.Writer, r *os.File) (string, string, status.S) {
	if _, err := w.Write([]byte("User Ident (e.g. admin): \n")); err != nil {
		return "", "", status.From(err)
	}
	br := bufio.NewReader(r)
	ident, err := br.ReadString('\n')
	if err != nil {
		return "", "", status.From(err)
	}

	if _, err := w.Write([]byte("Secret (e.g. 12345): \n")); err != nil {
		return "", "", status.From(err)
	}
	secret, err := terminal.ReadPassword(int(r.Fd()))
	if err != nil {
		return "", "", status.From(err)
	}

	return ident[:len(ident)-1], string(secret), nil
}
<|endoftext|>


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

// Package main contains the ctpu main program.
package main

import (
	"context"
	"fmt"
	"os"

	"flag"
	// context is used to cancel outstanding requests
	"github.com/google/subcommands"
	"github.com/tensorflow/tpu/tools/ctpu/commands"
	"github.com/tensorflow/tpu/tools/ctpu/config"
	"github.com/tensorflow/tpu/tools/ctpu/ctrl"
)

const version = "1.11"

var logRequests bool

func check(context string, err error) {
	if err != nil {
		fmt.Printf("Error encountered while %s: %s\n", context, err)
		os.Exit(1)
	}
}

func init() {
	flag.BoolVar(&logRequests, "log-http", false,
		`When debugging the ctpu tool, it can be helpful to see the full content of
        all HTTP request-response pairs. Set this flag to true in order to print out
        the requests.`)
}

func main() {
	cfg, err := config.FromEnv()
	check("creating configuration", err)
	cfg.SetFlags(flag.CommandLine)
	flag.Parse()

	ctx := context.Background()
	ctrls, err := ctrl.New(ctx, cfg, version, logRequests)
	check("creating API clients", err)

	subcommands.ImportantFlag("name")
	subcommands.ImportantFlag("project")
	subcommands.ImportantFlag("zone")

	subcommands.Register(commands.UpCommand(cfg, ctrls.TPU, ctrls.GCE, ctrls.ResourceManagement, ctrls.CLI), "")
	subcommands.Register(commands.PauseCommand(cfg, ctrls.TPU, ctrls.GCE), "")
	subcommands.Register(commands.DeleteCommand(cfg, ctrls.TPU, ctrls.GCE), "")
	subcommands.Register(commands.StatusCommand(cfg, ctrls.TPU, ctrls.GCE), "")
	subcommands.Register(commands.RestartCommand(cfg, ctrls.TPU), "")

	subcommands.Register(commands.AuthCommand(cfg, ctrls.ResourceManagement), "configuration")
	subcommands.Register(commands.ConfigCommand(cfg, ctrls.CLI), "configuration")
	subcommands.Register(commands.VersionCommand(version), "configuration")
	subcommands.Register(commands.ListCommand(cfg, ctrls.TPU, ctrls.GCE), "configuration")
	subcommands.Register(commands.TFVersionsCommand(cfg, ctrls.TPU), "configuration")
	subcommands.Register(commands.TPULocationsCommand(cfg, ctrls.TPU), "configuration")
	subcommands.Register(commands.QuotaCommand(cfg), "configuration")
	subcommands.Register(commands.TPUSizeCommand(cfg, ctrls.TPU), "configuration")

	subcommands.Register(commands.PauseCommandAlias(cfg, ctrls.TPU, ctrls.GCE), "aliases for other commands")
	subcommands.Register(commands.DeleteCommandAlias(cfg, ctrls.TPU, ctrls.GCE), "aliases for other commands")
	subcommands.Register(commands.StatusCommandAlias(cfg, ctrls.TPU, ctrls.GCE), "aliases for other commands")
	subcommands.Register(commands.ListCommandAlias(cfg, ctrls.TPU, ctrls.GCE), "aliases for other commands")
	subcommands.Register(commands.ConfigCommandAlias(cfg, ctrls.CLI), "aliases for other commands")

	subcommands.Register(subcommands.HelpCommand(), "usage")
	subcommands.Register(subcommands.FlagsCommand(), "usage")
	subcommands.Register(subcommands.CommandsCommand(), "usage")

	os.Exit(int(subcommands.Execute(ctx)))
}
<commit_msg>Version bump to fix tpu only configuration message.<commit_after>// Copyright 2017 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================

// Package main contains the ctpu main program.
package main

import (
	"context"
	"fmt"
	"os"

	"flag"
	// context is used to cancel outstanding requests
	"github.com/google/subcommands"
	"github.com/tensorflow/tpu/tools/ctpu/commands"
	"github.com/tensorflow/tpu/tools/ctpu/config"
	"github.com/tensorflow/tpu/tools/ctpu/ctrl"
)

const version = "1.12"

var logRequests bool

func check(context string, err error) {
	if err != nil {
		fmt.Printf("Error encountered while %s: %s\n", context, err)
		os.Exit(1)
	}
}

func init() {
	flag.BoolVar(&logRequests, "log-http", false,
		`When debugging the ctpu tool, it can be helpful to see the full content of
        all HTTP request-response pairs. Set this flag to true in order to print out
        the requests.`)
}

func main() {
	cfg, err := config.FromEnv()
	check("creating configuration", err)
	cfg.SetFlags(flag.CommandLine)
	flag.Parse()

	ctx := context.Background()
	ctrls, err := ctrl.New(ctx, cfg, version, logRequests)
	check("creating API clients", err)

	subcommands.ImportantFlag("name")
	subcommands.ImportantFlag("project")
	subcommands.ImportantFlag("zone")

	subcommands.Register(commands.UpCommand(cfg, ctrls.TPU, ctrls.GCE, ctrls.ResourceManagement, ctrls.CLI), "")
	subcommands.Register(commands.PauseCommand(cfg, ctrls.TPU, ctrls.GCE), "")
	subcommands.Register(commands.DeleteCommand(cfg, ctrls.TPU, ctrls.GCE), "")
	subcommands.Register(commands.StatusCommand(cfg, ctrls.TPU, ctrls.GCE), "")
	subcommands.Register(commands.RestartCommand(cfg, ctrls.TPU), "")

	subcommands.Register(commands.AuthCommand(cfg, ctrls.ResourceManagement), "configuration")
	subcommands.Register(commands.ConfigCommand(cfg, ctrls.CLI), "configuration")
	subcommands.Register(commands.VersionCommand(version), "configuration")
	subcommands.Register(commands.ListCommand(cfg, ctrls.TPU, ctrls.GCE), "configuration")
	subcommands.Register(commands.TFVersionsCommand(cfg, ctrls.TPU), "configuration")
	subcommands.Register(commands.TPULocationsCommand(cfg, ctrls.TPU), "configuration")
	subcommands.Register(commands.QuotaCommand(cfg), "configuration")
	subcommands.Register(commands.TPUSizeCommand(cfg, ctrls.TPU), "configuration")

	subcommands.Register(commands.PauseCommandAlias(cfg, ctrls.TPU, ctrls.GCE), "aliases for other commands")
	subcommands.Register(commands.DeleteCommandAlias(cfg, ctrls.TPU, ctrls.GCE), "aliases for other commands")
	subcommands.Register(commands.StatusCommandAlias(cfg, ctrls.TPU, ctrls.GCE), "aliases for other commands")
	subcommands.Register(commands.ListCommandAlias(cfg, ctrls.TPU, ctrls.GCE), "aliases for other commands")
	subcommands.Register(commands.ConfigCommandAlias(cfg, ctrls.CLI), "aliases for other commands")

	subcommands.Register(subcommands.HelpCommand(), "usage")
	subcommands.Register(subcommands.FlagsCommand(), "usage")
	subcommands.Register(subcommands.CommandsCommand(), "usage")

	os.Exit(int(subcommands.Execute(ctx)))
}
<|endoftext|>


<commit_before>// Package jsql provides glue between database/sql and json.
package jsql

import (
	"bytes"
	"database/sql"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"mime"
	"net/http"
	"regexp"

	"github.com/gorilla/mux"
)

const jsonContentType = "application/json;charset=UTF-8"

var reVars = regexp.MustCompile(`\${([^}]+)}`)

// PositionalQueryVars controls wether
// the generated query strings will use '?'
// or $1, $2... syntax.
var PositionalQueryVars = true

var Debug = false

// TODO(lvd) less simplistic parsing of query.  worry about quoting etc.
func rewrite(q string) (qq string, varnames []string) {
	idx := reVars.FindAllStringSubmatchIndex(q, -1)
	l := 0
	var pos = map[string]int{}
	var b bytes.Buffer
	for _, v := range idx {
		b.WriteString(q[l:v[0]])
		l = v[1]
		name := q[v[2]:v[3]]
		if PositionalQueryVars {
			if _, ok := pos[name]; !ok {
				varnames = append(varnames, name)
				pos[name] = len(varnames)
			}
			b.WriteString(fmt.Sprintf("$%d", pos[name]))
		} else {
			varnames = append(varnames, name)
			b.WriteString("?")
		}
	}
	b.WriteString(q[l:])
	return b.String(), varnames
}

// TODO:make Q send records on a channel, and provide json/bson as encoders
// then have the Handler switch returned encoding depending on http headers

// A queryfunc is the type of function generated by Q.
// It executes a query substituting arguments from the provided args map,
// and writes a JSON array of objects to the provided writer.
// It returns the number of array elements written and any error that may have occured.
//
// If w is an http.ResponseWriter (has a Header() http.Header method), then it will set
// the Content-Type header to application/json;charset=UTF-8 before the first write.
//
// If n == 0, nothing, not even the opening bracket, will have been written to w,
// and the header will not have been set, meaning it is not too late to call e.g.
//
//    http.Error(w, err.Error(), http.StatusInternalError)
//
// otherwise the error will have to be tacked on to the already produced output or
// reported out of band.  If any elements have been written the queryfunc will
// always generate a proper closing bracket.
type QueryFunc func(args map[string]interface{}, w io.Writer) (n int, err error)

// Q builds a function that can execute the query on a database given
// a map of parameters, and writes the result as json to a writer.
//
// The query must be of the form 'SELECT foo FROM bar WHERE x = ${varname}'
// that is, any query accepted by the driver that will provide db,
// but with the normal '$1, $2...' parameters replace by ${varname}.
//
// The generated function will extract the corresponding values from the args map
// and supply them to the db.Stmt.Execute method in the correct order.
// The query will be rewritten to use $1, $2 .. etc.
//
// The generated json will be an array of objects, one per row where
// each row object has keys named after the columns in the SQL query.
//
func Q(db *sql.DB, query string) (QueryFunc, error) {
	q, argn := rewrite(query)
	stmt, err := db.Prepare(q)
	if err != nil {
		return nil, err
	}
	return func(args map[string]interface{}, w io.Writer) (n int, err error) {
		var argv []interface{}
		for _, v := range argn {
			argv = append(argv, args[v])
		}
		rows, err := stmt.Query(argv...)
		if err != nil {
			return 0, err
		}
		defer rows.Close()
		retn, err := rows.Columns()
		if err != nil {
			return 0, err
		}

		var (
			retv  = make([]interface{}, len(retn))
			retvv = make([]interface{}, len(retn))
			retm  = make(map[string]interface{})
		)
		for i := range retv {
			retvv[i] = &retv[i]
			retm[retn[i]] = &retv[i]
		}
		for rows.Next() {
			if err := rows.Scan(retvv...); err != nil {
				return n, err
			}
			// convert []byte to strings, because json uuencodes []byte
			for i, v := range retv {
				if vv, ok := v.([]byte); ok {
					retv[i] = string(vv)
				}
			}

			b, err := json.Marshal(retm)
			if err != nil {
				return n, err
			}

			if n == 0 {
				if rw, ok := w.(http.ResponseWriter); ok {
					rw.Header().Set("Content-type", jsonContentType)
				}
				w.Write([]byte("[\n"))
				defer w.Write([]byte("\n]"))
			} else {
				w.Write([]byte(",\n"))
			}

			w.Write(b)
			n++
		}

		return n, rows.Err()
	}, nil
}

// Handler is a convenience wrapper around MkHandler
// that will die on error.
func Handler(db *sql.DB, query string) http.Handler {
	h, err := MkHandler(db, query)
	if err != nil {
		log.Panicf("MkHandler(%q): %v", query, err)
	}
	return h
}

// MkHandler produces a http.Handler that takes arguments from
// the request, and runs Q, producing a nice error message if needed.
//
// The arguments are taken from the request's formvalues or any
// input json object, depending on the requests content type, merged with
// the "github.com/gorilla/mux".Vars(r), which take precedence
// in case of name conflicts.
func MkHandler(db *sql.DB, query string) (http.Handler, error) {
	_, names := rewrite(query)

	qf, err := Q(db, query)
	if err != nil {
		return nil, err
	}

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		args := make(map[string]interface{})

		var jsonargs map[string]interface{}
		if r.Method == "POST" || r.Method == "PUT" {
			ct := r.Header.Get("Content-Type")
			if ct == "" {
				ct = "application/octet-stream"
			}
			ct, _, err = mime.ParseMediaType(ct)
			if ct == "application/json" {
				defer r.Body.Close()
				if err := json.NewDecoder(io.LimitReader(r.Body, 128<<10)).Decode(&jsonargs); err != nil {
					if err != io.EOF { // empty body is allowed
						http.Error(w, fmt.Sprintf("Can't decode json request: %v", err), http.StatusBadRequest)
						return
					}
				}
			}
		}
		muxargs := mux.Vars(r)
		// TODO harden against malicious input
		// for now we just rely on precedence: muxargs trum the others
		// but this still allows arbitrary values to end up in the query parameters
		for _, n := range names {
			if v, ok := muxargs[n]; ok {
				args[n] = v
				continue
			}
			if v := r.FormValue(n); v != "" {
				args[n] = v
				continue
			}
			if v, ok := jsonargs[n]; ok {
				args[n] = v
				continue
			}
		}
		if Debug {
			for k, v := range args {
				log.Printf("jsql arg %s = %#v", k, v)
			}
		}

		// if not all names set postgres $n deals with nil.
		n, err := qf(args, w)
		if err != nil {
			if n == 0 {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
			log.Print(r.Method, r.URL, ": ", err)
		}
	}), nil
}
<commit_msg>also parse arguments for DELETE<commit_after>// Package jsql provides glue between database/sql and json.
package jsql

import (
	"bytes"
	"database/sql"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"mime"
	"net/http"
	"regexp"

	"github.com/gorilla/mux"
)

const jsonContentType = "application/json;charset=UTF-8"

var reVars = regexp.MustCompile(`\${([^}]+)}`)

// PositionalQueryVars controls wether
// the generated query strings will use '?'
// or $1, $2... syntax.
var PositionalQueryVars = true

var Debug = false

// TODO(lvd) less simplistic parsing of query.  worry about quoting etc.
func rewrite(q string) (qq string, varnames []string) {
	idx := reVars.FindAllStringSubmatchIndex(q, -1)
	l := 0
	var pos = map[string]int{}
	var b bytes.Buffer
	for _, v := range idx {
		b.WriteString(q[l:v[0]])
		l = v[1]
		name := q[v[2]:v[3]]
		if PositionalQueryVars {
			if _, ok := pos[name]; !ok {
				varnames = append(varnames, name)
				pos[name] = len(varnames)
			}
			b.WriteString(fmt.Sprintf("$%d", pos[name]))
		} else {
			varnames = append(varnames, name)
			b.WriteString("?")
		}
	}
	b.WriteString(q[l:])
	return b.String(), varnames
}

// TODO:make Q send records on a channel, and provide json/bson as encoders
// then have the Handler switch returned encoding depending on http headers

// A queryfunc is the type of function generated by Q.
// It executes a query substituting arguments from the provided args map,
// and writes a JSON array of objects to the provided writer.
// It returns the number of array elements written and any error that may have occured.
//
// If w is an http.ResponseWriter (has a Header() http.Header method), then it will set
// the Content-Type header to application/json;charset=UTF-8 before the first write.
//
// If n == 0, nothing, not even the opening bracket, will have been written to w,
// and the header will not have been set, meaning it is not too late to call e.g.
//
//    http.Error(w, err.Error(), http.StatusInternalError)
//
// otherwise the error will have to be tacked on to the already produced output or
// reported out of band.  If any elements have been written the queryfunc will
// always generate a proper closing bracket.
type QueryFunc func(args map[string]interface{}, w io.Writer) (n int, err error)

// Q builds a function that can execute the query on a database given
// a map of parameters, and writes the result as json to a writer.
//
// The query must be of the form 'SELECT foo FROM bar WHERE x = ${varname}'
// that is, any query accepted by the driver that will provide db,
// but with the normal '$1, $2...' parameters replace by ${varname}.
//
// The generated function will extract the corresponding values from the args map
// and supply them to the db.Stmt.Execute method in the correct order.
// The query will be rewritten to use $1, $2 .. etc.
//
// The generated json will be an array of objects, one per row where
// each row object has keys named after the columns in the SQL query.
//
func Q(db *sql.DB, query string) (QueryFunc, error) {
	q, argn := rewrite(query)
	stmt, err := db.Prepare(q)
	if err != nil {
		return nil, err
	}
	return func(args map[string]interface{}, w io.Writer) (n int, err error) {
		var argv []interface{}
		for _, v := range argn {
			argv = append(argv, args[v])
		}
		rows, err := stmt.Query(argv...)
		if err != nil {
			return 0, err
		}
		defer rows.Close()
		retn, err := rows.Columns()
		if err != nil {
			return 0, err
		}

		var (
			retv  = make([]interface{}, len(retn))
			retvv = make([]interface{}, len(retn))
			retm  = make(map[string]interface{})
		)
		for i := range retv {
			retvv[i] = &retv[i]
			retm[retn[i]] = &retv[i]
		}
		for rows.Next() {
			if err := rows.Scan(retvv...); err != nil {
				return n, err
			}
			// convert []byte to strings, because json uuencodes []byte
			for i, v := range retv {
				if vv, ok := v.([]byte); ok {
					retv[i] = string(vv)
				}
			}

			b, err := json.Marshal(retm)
			if err != nil {
				return n, err
			}

			if n == 0 {
				if rw, ok := w.(http.ResponseWriter); ok {
					rw.Header().Set("Content-type", jsonContentType)
				}
				w.Write([]byte("[\n"))
				defer w.Write([]byte("\n]"))
			} else {
				w.Write([]byte(",\n"))
			}

			w.Write(b)
			n++
		}

		return n, rows.Err()
	}, nil
}

// Handler is a convenience wrapper around MkHandler
// that will die on error.
func Handler(db *sql.DB, query string) http.Handler {
	h, err := MkHandler(db, query)
	if err != nil {
		log.Panicf("MkHandler(%q): %v", query, err)
	}
	return h
}

// MkHandler produces a http.Handler that takes arguments from
// the request, and runs Q, producing a nice error message if needed.
//
// The arguments are taken from the request's formvalues or any
// input json object, depending on the requests content type, merged with
// the "github.com/gorilla/mux".Vars(r), which take precedence
// in case of name conflicts.
func MkHandler(db *sql.DB, query string) (http.Handler, error) {
	_, names := rewrite(query)

	qf, err := Q(db, query)
	if err != nil {
		return nil, err
	}

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		args := make(map[string]interface{})

		var jsonargs map[string]interface{}
		if r.Method == "POST" || r.Method == "PUT" || r.Method == "DELETE" {
			ct := r.Header.Get("Content-Type")
			if ct == "" {
				ct = "application/octet-stream"
			}
			ct, _, err = mime.ParseMediaType(ct)
			if ct == "application/json" {
				defer r.Body.Close()
				if err := json.NewDecoder(io.LimitReader(r.Body, 128<<10)).Decode(&jsonargs); err != nil {
					if err != io.EOF { // empty body is allowed
						http.Error(w, fmt.Sprintf("Can't decode json request: %v", err), http.StatusBadRequest)
						return
					}
				}
			}
		}
		muxargs := mux.Vars(r)
		// TODO harden against malicious input
		// for now we just rely on precedence: muxargs trum the others
		// but this still allows arbitrary values to end up in the query parameters
		for _, n := range names {
			if v, ok := muxargs[n]; ok {
				args[n] = v
				continue
			}
			if v := r.FormValue(n); v != "" {
				args[n] = v
				continue
			}
			if v, ok := jsonargs[n]; ok {
				args[n] = v
				continue
			}
		}
		if Debug {
			for k, v := range args {
				log.Printf("jsql arg %s = %#v", k, v)
			}
		}

		// if not all names set postgres $n deals with nil.
		n, err := qf(args, w)
		if err != nil {
			if n == 0 {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
			log.Print(r.Method, r.URL, ": ", err)
		}
	}), nil
}
<|endoftext|>


<commit_before>// Package juno implements the messaging spec of the Jupyter Messaging Protocol,
// just for you.
// http://jupyter-client.readthedocs.org/en/latest/messaging.html
package juno

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"os"

	zmq "github.com/pebbe/zmq4"
)

// MessageHeader is a Jupyter message header
// http://jupyter-client.readthedocs.org/en/latest/messaging.html
type MessageHeader struct {
	MessageID   string `json:"msg_id"`
	Username    string `json:"username"`
	Session     string `json:"session"`
	MessageType string `json:"msg_type"`
	Version     string `json:"version"`
}

// Message is a generic Jupyter message (not a wire message)
// http://jupyter-client.readthedocs.org/en/latest/messaging.html#general-message-format
type Message struct {
	Header       MessageHeader          `json:"header"`
	ParentHeader MessageHeader          `json:"parent_header"`
	Metadata     map[string]interface{} `json:"metadata"`
	Content      map[string]interface{} `json:"content"`
}

// MimeBundle is a collection of `mimetypes -> data`
// Example:
//     'text/html' -> '
Hey!
'
//     'image/png' -> 'R0lGODlhAQABAIAAAAAAAP///yH5BAEAAAAALAAAAAABAAEAAAIBRAA7'
type MimeBundle map[string]string

// ConnectionInfo represents the runtime connection data used by Jupyter kernels
type ConnectionInfo struct {
	IOPubPort       int    `json:"iopub_port"`
	StdinPort       int    `json:"stdin_port"`
	IP              string `json:"ip"`
	Transport       string `json:"transport"`
	HBPort          int    `json:"hb_port"`
	SignatureScheme string `json:"signature_scheme"`
	ShellPort       int    `json:"shell_port"`
	ControlPort     int    `json:"control_port"`
	Key             string `json:"key"`
}

// DELIMITER denotes the Jupyter multipart message
const DELIMITER = "<IDS|MSG>"

// ParseWireProtocol fills a Message with all the juicy Jupyter bits
func (m *Message) ParseWireProtocol(wireMessage [][]byte, key []byte) (err error) {
	var i int
	var el []byte

	// Wire protocol
	// http://jupyter-client.readthedocs.org/en/latest/messaging.html#the-wire-protocol
	/**
		[
	  		b'u-u-i-d',         # zmq identity(ies)
	  		b'<IDS|MSG>',       # delimiter
	  		b'baddad42',        # HMAC signature
			b'{header}',        # serialized header dict
			b'{parent_header}', # serialized parent header dict
			b'{metadata}',      # serialized metadata dict
			b'{content}',       # serialized content dict
			b'blob',            # extra raw data buffer(s)
	  		...
		]
	*/
	// Determine where the delimiter is
	for i, el = range wireMessage {
		if string(el) == DELIMITER {
			break // Found our delimiter
		}
	}

	if string(wireMessage[i]) != DELIMITER {
		return errors.New("Couldn't find delimeter")
	}

	// Extract the zmq identiti(es)
	//identities := wireMessage[:delimiterLocation]

	// If the message was signed
	if len(key) != 0 {
		// TODO: Programmatic selection of scheme
		mac := hmac.New(sha256.New, key)
		for _, msgpart := range wireMessage[i+2 : i+6] {
			mac.Write(msgpart)
		}
		signature := make([]byte, hex.DecodedLen(len(wireMessage[i+1])))
		hex.Decode(signature, wireMessage[i+1])
		if !hmac.Equal(mac.Sum(nil), signature) {
			return errors.New("Invalid signature")
		}
	}

	json.Unmarshal(wireMessage[i+2], &m.Header)
	json.Unmarshal(wireMessage[i+3], &m.ParentHeader)
	json.Unmarshal(wireMessage[i+4], &m.Metadata)
	json.Unmarshal(wireMessage[i+5], &m.Content)

	return nil
}

// OpenConnectionFile is a helper method that opens a connection file and reads
// it into a ConnectionInfo struct
func OpenConnectionFile(filename string) (ConnectionInfo, error) {
	var connInfo ConnectionInfo
	connFile, err := os.Open(filename)
	if err != nil {
		return connInfo, fmt.Errorf("Couldn't open connection file: %v", err)
	}

	jsonParser := json.NewDecoder(connFile)

	if err = jsonParser.Decode(&connInfo); err != nil {
		return connInfo, fmt.Errorf("Couldn't parse connection file: %v", err)
	}

	return connInfo, nil
}

// JupyterSocket is a zmq.Socket coupled with connection information
type JupyterSocket struct {
	ZMQSocket *zmq.Socket
	ConnInfo  ConnectionInfo
}

// ReadMessage reads a Jupyter Protocol Message
func (s *JupyterSocket) ReadMessage() (Message, error) {
	var message Message
	wireMessage, err := s.ZMQSocket.RecvMessageBytes(0)
	if err != nil {
		return message, fmt.Errorf("Error on receive: %v", err)
	}

	message.ParseWireProtocol(wireMessage, []byte(s.ConnInfo.Key))

	return message, nil
}

// Close shutdowns zmq sockets
func (s *JupyterSocket) Close() {
	s.ZMQSocket.Close()
}

// NewIOPubSocket creates a new IOPub socket (on SUB) with the connInfo given
func NewIOPubSocket(connInfo ConnectionInfo, subscribe string) (*JupyterSocket, error) {
	rawIOPubSocket, err := zmq.NewSocket(zmq.SUB)
	if err != nil {
		return nil, err
	}

	connectionString := fmt.Sprintf("%s://%s:%d",
		connInfo.Transport,
		connInfo.IP,
		connInfo.IOPubPort)

	rawIOPubSocket.Connect(connectionString)
	rawIOPubSocket.SetSubscribe("")

	iopub := &JupyterSocket{
		ZMQSocket: rawIOPubSocket,
		ConnInfo:  connInfo,
	}

	return iopub, nil
}
<commit_msg>Simpler assertion for finding the DELIMITER<commit_after>// Package juno implements the messaging spec of the Jupyter Messaging Protocol,
// just for you.
// http://jupyter-client.readthedocs.org/en/latest/messaging.html
package juno

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"os"

	zmq "github.com/pebbe/zmq4"
)

// MessageHeader is a Jupyter message header
// http://jupyter-client.readthedocs.org/en/latest/messaging.html
type MessageHeader struct {
	MessageID   string `json:"msg_id"`
	Username    string `json:"username"`
	Session     string `json:"session"`
	MessageType string `json:"msg_type"`
	Version     string `json:"version"`
}

// Message is a generic Jupyter message (not a wire message)
// http://jupyter-client.readthedocs.org/en/latest/messaging.html#general-message-format
type Message struct {
	Header       MessageHeader          `json:"header"`
	ParentHeader MessageHeader          `json:"parent_header"`
	Metadata     map[string]interface{} `json:"metadata"`
	Content      map[string]interface{} `json:"content"`
}

// MimeBundle is a collection of `mimetypes -> data`
// Example:
//     'text/html' -> '
Hey!
'
//     'image/png' -> 'R0lGODlhAQABAIAAAAAAAP///yH5BAEAAAAALAAAAAABAAEAAAIBRAA7'
type MimeBundle map[string]string

// ConnectionInfo represents the runtime connection data used by Jupyter kernels
type ConnectionInfo struct {
	IOPubPort       int    `json:"iopub_port"`
	StdinPort       int    `json:"stdin_port"`
	IP              string `json:"ip"`
	Transport       string `json:"transport"`
	HBPort          int    `json:"hb_port"`
	SignatureScheme string `json:"signature_scheme"`
	ShellPort       int    `json:"shell_port"`
	ControlPort     int    `json:"control_port"`
	Key             string `json:"key"`
}

// DELIMITER denotes the Jupyter multipart message
const DELIMITER = "<IDS|MSG>"

// ParseWireProtocol fills a Message with all the juicy Jupyter bits
func (m *Message) ParseWireProtocol(wireMessage [][]byte, key []byte) (err error) {
	var i int
	var el []byte

	// Wire protocol
	// http://jupyter-client.readthedocs.org/en/latest/messaging.html#the-wire-protocol
	/**
		[
	  		b'u-u-i-d',         # zmq identity(ies)
	  		b'<IDS|MSG>',       # delimiter
	  		b'baddad42',        # HMAC signature
			b'{header}',        # serialized header dict
			b'{parent_header}', # serialized parent header dict
			b'{metadata}',      # serialized metadata dict
			b'{content}',       # serialized content dict
			b'blob',            # extra raw data buffer(s)
	  		...
		]
	*/
	// Determine where the delimiter is
	for i, el = range wireMessage {
		if string(el) == DELIMITER {
			break // Found our delimiter
		}
	}

	if i >= len(wireMessage) {
		return errors.New("Couldn't find delimeter")
	}

	// Extract the zmq identiti(es)
	//identities := wireMessage[:delimiterLocation]

	// If the message was signed
	if len(key) != 0 {
		// TODO: Programmatic selection of scheme
		mac := hmac.New(sha256.New, key)
		for _, msgpart := range wireMessage[i+2 : i+6] {
			mac.Write(msgpart)
		}
		signature := make([]byte, hex.DecodedLen(len(wireMessage[i+1])))
		hex.Decode(signature, wireMessage[i+1])
		if !hmac.Equal(mac.Sum(nil), signature) {
			return errors.New("Invalid signature")
		}
	}

	json.Unmarshal(wireMessage[i+2], &m.Header)
	json.Unmarshal(wireMessage[i+3], &m.ParentHeader)
	json.Unmarshal(wireMessage[i+4], &m.Metadata)
	json.Unmarshal(wireMessage[i+5], &m.Content)

	return nil
}

// OpenConnectionFile is a helper method that opens a connection file and reads
// it into a ConnectionInfo struct
func OpenConnectionFile(filename string) (ConnectionInfo, error) {
	var connInfo ConnectionInfo
	connFile, err := os.Open(filename)
	if err != nil {
		return connInfo, fmt.Errorf("Couldn't open connection file: %v", err)
	}

	jsonParser := json.NewDecoder(connFile)

	if err = jsonParser.Decode(&connInfo); err != nil {
		return connInfo, fmt.Errorf("Couldn't parse connection file: %v", err)
	}

	return connInfo, nil
}

// JupyterSocket is a zmq.Socket coupled with connection information
type JupyterSocket struct {
	ZMQSocket *zmq.Socket
	ConnInfo  ConnectionInfo
}

// ReadMessage reads a Jupyter Protocol Message
func (s *JupyterSocket) ReadMessage() (Message, error) {
	var message Message
	wireMessage, err := s.ZMQSocket.RecvMessageBytes(0)
	if err != nil {
		return message, fmt.Errorf("Error on receive: %v", err)
	}

	message.ParseWireProtocol(wireMessage, []byte(s.ConnInfo.Key))

	return message, nil
}

// Close shutdowns zmq sockets
func (s *JupyterSocket) Close() {
	s.ZMQSocket.Close()
}

// NewIOPubSocket creates a new IOPub socket (on SUB) with the connInfo given
func NewIOPubSocket(connInfo ConnectionInfo, subscribe string) (*JupyterSocket, error) {
	rawIOPubSocket, err := zmq.NewSocket(zmq.SUB)
	if err != nil {
		return nil, err
	}

	connectionString := fmt.Sprintf("%s://%s:%d",
		connInfo.Transport,
		connInfo.IP,
		connInfo.IOPubPort)

	rawIOPubSocket.Connect(connectionString)
	rawIOPubSocket.SetSubscribe("")

	iopub := &JupyterSocket{
		ZMQSocket: rawIOPubSocket,
		ConnInfo:  connInfo,
	}

	return iopub, nil
}
<|endoftext|>


<commit_before>package kana

import (
	"io/ioutil"
	"strings"
)

type Kana struct {
	kanaToRomajiTrie     *Trie
	romajiToHiraganaTrie *Trie
	romajiToKatakanaTrie *Trie
}

func newKana() *Kana {
	/*
	   Build a trie for efficient retrieval of entries
	*/
	kana := &Kana{newTrie(), newTrie(), newTrie()}
	kana.initialize()
	return kana
}

func (k *Kana) initialize() {
	/*
		Build the Hiragana + Katakana trie.

		Because there is no overlap between the hiragana and katakana sets,
		they both use the same trie without conflict. Nice bonus!
	*/
	filenames := []string{"hiragana.in", "katakana.in"}
	for _, filename := range filenames {
		content, err := ioutil.ReadFile(filename)
		if err != nil {
			//Do something
		}
		rows := strings.Split(string(content), "\n")
		colNames := strings.Split(string(rows[0]), "\t")[1:]
		for _, row := range rows[1:] {
			cols := strings.Split(string(row), "\t")
			rowName := cols[0]
			for i, kana := range cols[1:] {
				value := rowName + colNames[i]
				kanas := strings.Split(kana, "/")
				for _, singleKana := range kanas {
					if singleKana != "" {
						// add to tries
						k.kanaToRomajiTrie.insert(singleKana, value)
						if filename == "hiragana.in" {
							k.romajiToHiraganaTrie.insert(value, singleKana)
						} else if filename == "katakana.in" {
							k.romajiToKatakanaTrie.insert(value, singleKana)
						}
					}
				}
			}
		}
	}
}

func (k Kana) kana_to_romaji(kana string) (romaji string) {
	romaji = k.kanaToRomajiTrie.convert(kana)

	// do some post-processing for the tsu and stripe characters
	// maybe a bit of a hacky solution - how can we improve?
	// (they act more like punctuation)
	tsus := []string{"っ", "ッ"}
	for _, tsu := range tsus {
		for i := strings.Index(romaji, tsu); i > -1; i = strings.Index(romaji, tsu) {
			rune_romaji := []rune(romaji)
			if len(rune_romaji) > i+2 {
				// TODO: should check if following letter is consonant
				followingLetter := string(rune_romaji[i+1 : i+2])
				romaji = strings.Replace(romaji, tsu, followingLetter, 1)
			} else {
				romaji = strings.Replace(romaji, tsu, "", 1)
			}
		}
	}

	line := "ー"
	for i := strings.Index(romaji, line); i > -1; i = strings.Index(romaji, line) {
		rune_romaji := []rune(romaji)
		if i > 0 {
			// TODO: should check if following letter is consonant
			previousLetter := string(rune_romaji[i-1 : i])
			romaji = strings.Replace(romaji, line, previousLetter, 1)
		} else {
			romaji = strings.Replace(romaji, line, "", 1)
		}
	}

	return romaji
}

func (k Kana) romaji_to_hiragana(romaji string) (hiragana string) {
	hiragana = k.romajiToHiraganaTrie.convert(romaji)
	return hiragana
}

func (k Kana) romaji_to_katakana(romaji string) (katakana string) {
	katakana = k.romajiToKatakanaTrie.convert(romaji)
	return katakana
}
<commit_msg>convert dashes to long dashes<commit_after>package kana

import (
	"io/ioutil"
	"strings"
)

type Kana struct {
	kanaToRomajiTrie     *Trie
	romajiToHiraganaTrie *Trie
	romajiToKatakanaTrie *Trie
}

func newKana() *Kana {
	/*
	   Build a trie for efficient retrieval of entries
	*/
	kana := &Kana{newTrie(), newTrie(), newTrie()}
	kana.initialize()
	return kana
}

func (k *Kana) initialize() {
	/*
		Build the Hiragana + Katakana trie.

		Because there is no overlap between the hiragana and katakana sets,
		they both use the same trie without conflict. Nice bonus!
	*/
	filenames := []string{"hiragana.in", "katakana.in"}
	for _, filename := range filenames {
		content, err := ioutil.ReadFile(filename)
		if err != nil {
			//Do something
		}
		rows := strings.Split(string(content), "\n")
		colNames := strings.Split(string(rows[0]), "\t")[1:]
		for _, row := range rows[1:] {
			cols := strings.Split(string(row), "\t")
			rowName := cols[0]
			for i, kana := range cols[1:] {
				value := rowName + colNames[i]
				kanas := strings.Split(kana, "/")
				for _, singleKana := range kanas {
					if singleKana != "" {
						// add to tries
						k.kanaToRomajiTrie.insert(singleKana, value)
						if filename == "hiragana.in" {
							k.romajiToHiraganaTrie.insert(value, singleKana)
						} else if filename == "katakana.in" {
							k.romajiToKatakanaTrie.insert(value, singleKana)
						}
					}
				}
			}
		}
	}
}

func (k Kana) kana_to_romaji(kana string) (romaji string) {
	romaji = k.kanaToRomajiTrie.convert(kana)

	// do some post-processing for the tsu and stripe characters
	// maybe a bit of a hacky solution - how can we improve?
	// (they act more like punctuation)
	tsus := []string{"っ", "ッ"}
	for _, tsu := range tsus {
		for i := strings.Index(romaji, tsu); i > -1; i = strings.Index(romaji, tsu) {
			rune_romaji := []rune(romaji)
			if len(rune_romaji) > i+2 {
				// TODO: should check if following letter is consonant
				followingLetter := string(rune_romaji[i+1 : i+2])
				romaji = strings.Replace(romaji, tsu, followingLetter, 1)
			} else {
				romaji = strings.Replace(romaji, tsu, "", 1)
			}
		}
	}

	line := "ー"
	for i := strings.Index(romaji, line); i > -1; i = strings.Index(romaji, line) {
		rune_romaji := []rune(romaji)
		if i > 0 {
			// TODO: should check if following letter is consonant
			previousLetter := string(rune_romaji[i-1 : i])
			romaji = strings.Replace(romaji, line, previousLetter, 1)
		} else {
			romaji = strings.Replace(romaji, line, "", 1)
		}
	}

	return romaji
}

func (k Kana) romaji_to_hiragana(romaji string) (hiragana string) {
	romaji = strings.Replace(romaji, "-", "ー", -1)
	hiragana = k.romajiToHiraganaTrie.convert(romaji)
	return hiragana
}

func (k Kana) romaji_to_katakana(romaji string) (katakana string) {
	romaji = strings.Replace(romaji, "-", "ー", -1)
	katakana = k.romajiToKatakanaTrie.convert(romaji)
	return katakana
}
<|endoftext|>


<commit_before>/* Copyright 2016-2017 Vector Creations Ltd
 *
 * 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 gomatrixserverlib

import (
	"bufio"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"strings"
	"time"

	"github.com/pkg/errors"
)

// ServerKeys are the ed25519 signing keys published by a matrix server.
// Contains SHA256 fingerprints of the TLS X509 certificates used by the server.
type ServerKeys struct {
	// Copy of the raw JSON for signature checking.
	Raw []byte
	// The decoded JSON fields.
	ServerKeyFields
}

// A TLSFingerprint is a SHA256 hash of an X509 certificate.
type TLSFingerprint struct {
	SHA256 Base64String `json:"sha256"`
}

// A VerifyKey is a ed25519 public key for a server.
type VerifyKey struct {
	// The public key.
	Key Base64String `json:"key"`
}

// An OldVerifyKey is an old ed25519 public key that is no longer valid.
type OldVerifyKey struct {
	VerifyKey
	// When this key stopped being valid for event signing in milliseconds.
	ExpiredTS Timestamp `json:"expired_ts"`
}

// ServerKeyFields are the parsed JSON contents of the ed25519 signing keys published by a matrix server.
type ServerKeyFields struct {
	// The name of the server
	ServerName ServerName `json:"server_name"`
	// List of SHA256 fingerprints of X509 certificates used by this server.
	TLSFingerprints []TLSFingerprint `json:"tls_fingerprints"`
	// The current signing keys in use on this server.
	// The keys of the map are the IDs of the keys.
	// These are valid while this response is valid.
	VerifyKeys map[KeyID]VerifyKey `json:"verify_keys"`
	// When this result is valid until in milliseconds.
	ValidUntilTS Timestamp `json:"valid_until_ts"`
	// Old keys that are now only valid for checking historic events.
	// The keys of the map are the IDs of the keys.
	OldVerifyKeys map[KeyID]OldVerifyKey `json:"old_verify_keys"`
}

// UnmarshalJSON implements json.Unmarshaler
func (keys *ServerKeys) UnmarshalJSON(data []byte) error {
	keys.Raw = data
	return json.Unmarshal(data, &keys.ServerKeyFields)
}

// MarshalJSON implements json.Marshaler
func (keys ServerKeys) MarshalJSON() ([]byte, error) {
	// We already have a copy of the serialised JSON for the keys so we can return that directly.
	return keys.Raw, nil
}

// PublicKey returns a public key with the given ID valid at the given TS or nil if no such key exists.
func (keys ServerKeys) PublicKey(keyID KeyID, atTS Timestamp) []byte {
	if currentKey, ok := keys.VerifyKeys[keyID]; ok && (atTS <= keys.ValidUntilTS) {
		return currentKey.Key
	}
	if oldKey, ok := keys.OldVerifyKeys[keyID]; ok && (atTS <= oldKey.ExpiredTS) {
		return oldKey.Key
	}
	return nil
}

// FetchKeysDirect fetches the matrix keys directly from the given address.
// Optionally sets a SNI header if ``sni`` is not empty.
// Returns the server keys and the state of the TLS connection used to retrieve them.
func FetchKeysDirect(serverName ServerName, addr, sni string) (*ServerKeys, *tls.ConnectionState, error) {
	// Create a TLS connection.
	tcpconn, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, nil, err
	}
	defer tcpconn.Close() // nolint: errcheck
	tlsconn := tls.Client(tcpconn, &tls.Config{
		ServerName: sni,

		// This must be specified even though the TLS library will ignore it.
		InsecureSkipVerify: true, // nolint: gas
	})
	if err = tlsconn.Handshake(); err != nil {
		return nil, nil, err
	}
	connectionState := tlsconn.ConnectionState()

	// Write a GET /_matrix/key/v2/server down the connection.
	requestURL := "matrix://" + string(serverName) + "/_matrix/key/v2/server"
	request, err := http.NewRequest("GET", requestURL, nil)
	if err != nil {
		return nil, nil, err
	}
	request.Header.Set("Connection", "close")
	if err = request.Write(tlsconn); err != nil {
		return nil, nil, err
	}

	// Read the 200 OK from the server.
	response, err := http.ReadResponse(bufio.NewReader(tlsconn), request)
	if response != nil {
		defer response.Body.Close() // nolint: errcheck
	}
	if err != nil {
		return nil, nil, err
	}
	if response.StatusCode != http.StatusOK {
		return nil, nil, fmt.Errorf("Non-200 response %d from remote server", response.StatusCode)
	}
	var keys ServerKeys
	if err = json.NewDecoder(response.Body).Decode(&keys); err != nil {
		return nil, nil, errors.Wrap(err, "Unable to decode JSON from remote server")
	}
	return &keys, &connectionState, nil
}

// Ed25519Checks are the checks that are applied to Ed25519 keys in ServerKey responses.
type Ed25519Checks struct {
	ValidEd25519      bool // The verify key is valid Ed25519 keys.
	MatchingSignature bool // The verify key has a valid signature.
}

// KeyChecks are the checks that should be applied to ServerKey responses.
type KeyChecks struct {
	AllChecksOK        bool                    // Did all the checks pass?
	MatchingServerName bool                    // Does the server name match what was requested.
	FutureValidUntilTS bool                    // The valid until TS is in the future.
	HasEd25519Key      bool                    // The server has at least one ed25519 key.
	AllEd25519ChecksOK *bool                   // All the Ed25519 checks are ok. or null if there weren't any to check.
	Ed25519Checks      map[KeyID]Ed25519Checks // Checks for Ed25519 keys.
}

// CheckKeys checks the keys returned from a server to make sure they are valid.
// If the checks pass then also return a map of key_id to Ed25519 public key
func CheckKeys(
	serverName ServerName,
	now time.Time,
	keys ServerKeys,
) (
	checks KeyChecks,
	ed25519Keys map[KeyID]Base64String,
) {
	checks.MatchingServerName = serverName == keys.ServerName
	checks.FutureValidUntilTS = keys.ValidUntilTS.Time().After(now)
	checks.AllChecksOK = checks.MatchingServerName && checks.FutureValidUntilTS

	ed25519Keys = checkVerifyKeys(keys, &checks)

	if !checks.AllChecksOK {
		ed25519Keys = nil
	}
	return
}

func checkVerifyKeys(keys ServerKeys, checks *KeyChecks) map[KeyID]Base64String {
	allEd25519ChecksOK := true
	checks.Ed25519Checks = map[KeyID]Ed25519Checks{}
	verifyKeys := map[KeyID]Base64String{}
	for keyID, keyData := range keys.VerifyKeys {
		algorithm := strings.SplitN(string(keyID), ":", 2)[0]
		publicKey := keyData.Key
		if algorithm == "ed25519" {
			checks.HasEd25519Key = true
			checks.AllEd25519ChecksOK = &allEd25519ChecksOK
			entry := Ed25519Checks{
				ValidEd25519: len(publicKey) == 32,
			}
			if entry.ValidEd25519 {
				err := VerifyJSON(string(keys.ServerName), keyID, []byte(publicKey), keys.Raw)
				entry.MatchingSignature = err == nil
			}
			checks.Ed25519Checks[keyID] = entry
			if entry.MatchingSignature {
				verifyKeys[keyID] = publicKey
			} else {
				allEd25519ChecksOK = false
			}
		}
	}
	if checks.AllChecksOK {
		checks.AllChecksOK = checks.HasEd25519Key && allEd25519ChecksOK
	}
	return verifyKeys
}
<commit_msg>Move FetchKeysDirect to federation tester (#126)<commit_after>/* Copyright 2016-2017 Vector Creations Ltd
 *
 * 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 gomatrixserverlib

import (
	"encoding/json"
	"strings"
	"time"
)

// ServerKeys are the ed25519 signing keys published by a matrix server.
// Contains SHA256 fingerprints of the TLS X509 certificates used by the server.
type ServerKeys struct {
	// Copy of the raw JSON for signature checking.
	Raw []byte
	// The decoded JSON fields.
	ServerKeyFields
}

// A TLSFingerprint is a SHA256 hash of an X509 certificate.
type TLSFingerprint struct {
	SHA256 Base64String `json:"sha256"`
}

// A VerifyKey is a ed25519 public key for a server.
type VerifyKey struct {
	// The public key.
	Key Base64String `json:"key"`
}

// An OldVerifyKey is an old ed25519 public key that is no longer valid.
type OldVerifyKey struct {
	VerifyKey
	// When this key stopped being valid for event signing in milliseconds.
	ExpiredTS Timestamp `json:"expired_ts"`
}

// ServerKeyFields are the parsed JSON contents of the ed25519 signing keys published by a matrix server.
type ServerKeyFields struct {
	// The name of the server
	ServerName ServerName `json:"server_name"`
	// List of SHA256 fingerprints of X509 certificates used by this server.
	TLSFingerprints []TLSFingerprint `json:"tls_fingerprints"`
	// The current signing keys in use on this server.
	// The keys of the map are the IDs of the keys.
	// These are valid while this response is valid.
	VerifyKeys map[KeyID]VerifyKey `json:"verify_keys"`
	// When this result is valid until in milliseconds.
	ValidUntilTS Timestamp `json:"valid_until_ts"`
	// Old keys that are now only valid for checking historic events.
	// The keys of the map are the IDs of the keys.
	OldVerifyKeys map[KeyID]OldVerifyKey `json:"old_verify_keys"`
}

// UnmarshalJSON implements json.Unmarshaler
func (keys *ServerKeys) UnmarshalJSON(data []byte) error {
	keys.Raw = data
	return json.Unmarshal(data, &keys.ServerKeyFields)
}

// MarshalJSON implements json.Marshaler
func (keys ServerKeys) MarshalJSON() ([]byte, error) {
	// We already have a copy of the serialised JSON for the keys so we can return that directly.
	return keys.Raw, nil
}

// PublicKey returns a public key with the given ID valid at the given TS or nil if no such key exists.
func (keys ServerKeys) PublicKey(keyID KeyID, atTS Timestamp) []byte {
	if currentKey, ok := keys.VerifyKeys[keyID]; ok && (atTS <= keys.ValidUntilTS) {
		return currentKey.Key
	}
	if oldKey, ok := keys.OldVerifyKeys[keyID]; ok && (atTS <= oldKey.ExpiredTS) {
		return oldKey.Key
	}
	return nil
}

// Ed25519Checks are the checks that are applied to Ed25519 keys in ServerKey responses.
type Ed25519Checks struct {
	ValidEd25519      bool // The verify key is valid Ed25519 keys.
	MatchingSignature bool // The verify key has a valid signature.
}

// KeyChecks are the checks that should be applied to ServerKey responses.
type KeyChecks struct {
	AllChecksOK        bool                    // Did all the checks pass?
	MatchingServerName bool                    // Does the server name match what was requested.
	FutureValidUntilTS bool                    // The valid until TS is in the future.
	HasEd25519Key      bool                    // The server has at least one ed25519 key.
	AllEd25519ChecksOK *bool                   // All the Ed25519 checks are ok. or null if there weren't any to check.
	Ed25519Checks      map[KeyID]Ed25519Checks // Checks for Ed25519 keys.
}

// CheckKeys checks the keys returned from a server to make sure they are valid.
// If the checks pass then also return a map of key_id to Ed25519 public key
func CheckKeys(
	serverName ServerName,
	now time.Time,
	keys ServerKeys,
) (
	checks KeyChecks,
	ed25519Keys map[KeyID]Base64String,
) {
	checks.MatchingServerName = serverName == keys.ServerName
	checks.FutureValidUntilTS = keys.ValidUntilTS.Time().After(now)
	checks.AllChecksOK = checks.MatchingServerName && checks.FutureValidUntilTS

	ed25519Keys = checkVerifyKeys(keys, &checks)

	if !checks.AllChecksOK {
		ed25519Keys = nil
	}
	return
}

func checkVerifyKeys(keys ServerKeys, checks *KeyChecks) map[KeyID]Base64String {
	allEd25519ChecksOK := true
	checks.Ed25519Checks = map[KeyID]Ed25519Checks{}
	verifyKeys := map[KeyID]Base64String{}
	for keyID, keyData := range keys.VerifyKeys {
		algorithm := strings.SplitN(string(keyID), ":", 2)[0]
		publicKey := keyData.Key
		if algorithm == "ed25519" {
			checks.HasEd25519Key = true
			checks.AllEd25519ChecksOK = &allEd25519ChecksOK
			entry := Ed25519Checks{
				ValidEd25519: len(publicKey) == 32,
			}
			if entry.ValidEd25519 {
				err := VerifyJSON(string(keys.ServerName), keyID, []byte(publicKey), keys.Raw)
				entry.MatchingSignature = err == nil
			}
			checks.Ed25519Checks[keyID] = entry
			if entry.MatchingSignature {
				verifyKeys[keyID] = publicKey
			} else {
				allEd25519ChecksOK = false
			}
		}
	}
	if checks.AllChecksOK {
		checks.AllChecksOK = checks.HasEd25519Key && allEd25519ChecksOK
	}
	return verifyKeys
}
<|endoftext|>


<commit_before>package main

import "fmt"
import "strings"
import "errors"

func processGokContent(code string) (string, string, error) {
    var gofile string = "%s\nfunc %s(gok *Gok){%s\n}"
    funcName := fmt.Sprintf("Render%s", genRandName())
    imports, r, err := buildImports(code)
    if err != nil {
        return "", "", err
    }
    goCode, err := buildGoCode(code[r:])
    if err != nil {
        return "", "", err
    }
    final := fmt.Sprintf(gofile, imports, funcName, goCode)
    return final, funcName, nil
}

func buildImports(code string) (string, int, error) {
    p := "<?goUse"
    pe := "?>"
	plen := len(p)
	indx := strings.Index(code, p)
	if indx == -1 {
		return "package main\n", 0, nil
	}
	indxEnd := strings.Index(code, pe)
	if indxEnd == -1 {
		return "", -1, errors.New("unknown code pattern")
	}
	if indxEnd == (indx + plen) {
		return "package main\n", 0, nil
	}
	imports := strings.Split(code[(indx+plen):indxEnd], "\n")
	for i := range imports {
		imports[i] = strings.TrimSpace(imports[i])
	}
	return "package main\n" + strings.Join(imports, "\n"), (indxEnd + 2), nil
}

func buildGoCode(code string) (string, error) {
	echoFunc := "\ngok.Echo(\"%s\");"
	p := "<?go "
	pe := "?>"
	codeLen := len(code)
	result := make([]byte, 0, codeLen*3)
	for last := 0; last < codeLen; {
		slice := code[last:]
		indx := strings.Index(slice, p)
		if indx == -1 {
			if len(slice) > 0 {
				echo := fmt.Sprintf(echoFunc, strToCStr(slice))
				result = append(result, echo...)
			}
			break
		}
		if indx != 0 {
			echo := fmt.Sprintf(echoFunc, strToCStr(slice[0:indx]))
			result = append(result, echo...)
		}
		indxEnd := strings.Index(slice, pe)
		if indxEnd == -1 {
			return "", errors.New("unknown code pattern")
		}
		if indxEnd == (indx + 5) {
			last += indxEnd + 2
			continue
		}
		goCode := strings.TrimSpace(slice[indx+5 : indxEnd])
		if len(goCode) != 0 {
			for _, v := range strings.Split(goCode, "\n") {
				result = append(result, '\n')
				result = append(result, strings.TrimSpace(v)...)
			}
		}
		last += indxEnd + 2
	}
	return string(result), nil
}

func strToCStr(str string) string {
	result := strings.Replace(str, "\n", "\\n", -1)
	result = strings.Replace(result, "\t", "\\t", -1)
	result = strings.Replace(result, "\r", "\\r", -1)
	result = strings.Replace(result, "\v", "\\v", -1)
	result = strings.Replace(result, "\f", "\\f", -1)
	result = strings.Replace(result, "\"", "\\\"", -1)
	return result
}
<commit_msg>standard<commit_after>package main

import "fmt"
import "strings"
import "errors"

func processGokContent(code string) (string, string, error) {
    var gofile string = "%s\nfunc %s(gok *Gok){%s\n}"
    funcName := fmt.Sprintf("Render%s", genRandName())
    imports, r, err := buildImports(code)
    if err != nil {
        return "", "", err
    }
    goCode, err := buildGoCode(code[r:])
    if err != nil {
        return "", "", err
    }
    final := fmt.Sprintf(gofile, imports, funcName, goCode)
    return final, funcName, nil
}

func buildImports(code string) (string, int, error) {
    p := "<?goUse"
    pe := "?>"
    plen := len(p)
    indx := strings.Index(code, p)
    if indx == -1 {
        return "package main\n", 0, nil
    }
    indxEnd := strings.Index(code, pe)
    if indxEnd == -1 {
        return "", -1, errors.New("unknown code pattern")
    }
    if indxEnd == (indx + plen) {
        return "package main\n", 0, nil
    }
    imports := strings.Split(code[(indx+plen):indxEnd], "\n")
    for i := range imports {
        imports[i] = strings.TrimSpace(imports[i])
    }
    return "package main\n" + strings.Join(imports, "\n"), (indxEnd + 2), nil
}

func buildGoCode(code string) (string, error) {
    echoFunc := "\ngok.Echo(\"%s\");"
    p := "<?go "
    pe := "?>"
    codeLen := len(code)
    result := make([]byte, 0, codeLen*3)
    for last := 0; last < codeLen; {
        slice := code[last:]
        indx := strings.Index(slice, p)
        if indx == -1 {
            if len(slice) > 0 {
                echo := fmt.Sprintf(echoFunc, strToCStr(slice))
                result = append(result, echo...)
            }
            break
        }
        if indx != 0 {
            echo := fmt.Sprintf(echoFunc, strToCStr(slice[0:indx]))
            result = append(result, echo...)
        }
        indxEnd := strings.Index(slice, pe)
        if indxEnd == -1 {
            return "", errors.New("unknown code pattern")
        }
        if indxEnd == (indx + 5) {
            last += indxEnd + 2
            continue;
        }
        goCode := strings.TrimSpace(slice[indx+5 : indxEnd])
        if len(goCode) != 0 {
            for _, v := range strings.Split(goCode, "\n") {
                result = append(result, '\n')
                result = append(result, strings.TrimSpace(v)...)
            }
        }
        last += indxEnd + 2
    }
    return string(result), nil
}

func strToCStr(str string) string {
    result := strings.Replace(str, "\n", "\\n", -1)
    result = strings.Replace(result, "\t", "\\t", -1)
    result = strings.Replace(result, "\r", "\\r", -1)
    result = strings.Replace(result, "\v", "\\v", -1)
    result = strings.Replace(result, "\f", "\\f", -1)
    result = strings.Replace(result, "\"", "\\\"", -1)
    return result
}
<|endoftext|>


<commit_before>package edc

// Kind - struct for kind
type Kind struct {
	ID        int64  `sql:"id"         json:"id"         form:"id"         query:"id"`
	Name      string `sql:"name"       json:"name"       form:"name"       query:"name"`
	ShortName string `sql:"short_name" json:"short_name" form:"short_name" query:"short_name"`
	Note      string `sql:"note,null"  json:"note"       form:"note"       query:"note"`
	CreatedAt string `sql:"created_at" json:"-"`
	UpdatedAt string `sql:"updated_at" json:"-"`
}

// KindList - struct for kind list
type KindList struct {
	ID        int64  `sql:"id"         json:"id"         form:"id"         query:"id"`
	Name      string `sql:"name"       json:"name"       form:"name"       query:"name"`
	ShortName string `sql:"short_name" json:"short_name" form:"short_name" query:"short_name"`
	Note      string `sql:"note,null"  json:"note"       form:"note"       query:"note"`
}

// GetKind - get one kind by id
func (e *Edb) GetKind(id int64) (Kind, error) {
	var kind Kind
	if id == 0 {
		return kind, nil
	}
	err := e.db.Model(&kind).
		Where("id = ?", id).
		Select()
	if err != nil {
		errmsg("GetKind select", err)
	}
	return kind, err
}

// GetKindList - get kind for list by id
func (e *Edb) GetKindList(id int64) (KindList, error) {
	var kind KindList
	err := e.db.Model(&Kind{}).
		Column("id", "name", "short_name", "note").
		Where("id = ?").
		Select(&kind)
	if err != nil {
		errmsg("GetKindList select", err)
	}
	return kind, err
}

// GetKindListAll - get all kind for list
func (e *Edb) GetKindListAll() ([]KindList, error) {
	var kinds []KindList
	err := e.db.Model(&Kind{}).
		Column("id", "name", "short_name", "note").
		Order("name ASC").
		Select(&kinds)
	if err != nil {
		errmsg("GetKindListAll select", err)
	}
	return kinds, err
}

// GetKindSelectAll - get all kind for select
func (e *Edb) GetKindSelectAll() ([]SelectItem, error) {
	var kinds []SelectItem
	err := e.db.Model(&Kind{}).
		Column("id", "name").
		Order("name ASC").
		Select(&kinds)
	if err != nil {
		errmsg("GetKindSelectAll select", err)
	}
	return kinds, err
}

// CreateKind - create new kind
func (e *Edb) CreateKind(kind Kind) (int64, error) {
	err := e.db.Insert(&kind)
	if err != nil {
		errmsg("CreateKind insert", err)
	}
	return kind.ID, nil
}

// UpdateKind - save kind changes
func (e *Edb) UpdateKind(kind Kind) error {
	err := e.db.Update(&kind)
	if err != nil {
		errmsg("UpdateKind update", err)
	}
	return err
}

// DeleteKind - delete kind by id
func (e *Edb) DeleteKind(id int64) error {
	if id == 0 {
		return nil
	}
	_, err := e.db.Model(&Kind{}).
		Where("id = ?", id).
		Delete()
	if err != nil {
		errmsg("DeleteKind delete", err)
	}
	return err
}

func (e *Edb) kindCreateTable() error {
	str := `
		CREATE TABLE IF NOT EXISTS
			kinds (
				id bigserial primary key,
				name text,
				short_name text,
				note text,
				created_at TIMESTAMP without time zone,
				updated_at TIMESTAMP without time zone default now(),
				UNIQUE(name)
			)
	`
	_, err := e.db.Exec(str)
	if err != nil {
		errmsg("kindCreateTable exec", err)
	}
	return err
}
<commit_msg>fix GetKindList<commit_after>package edc

// Kind - struct for kind
type Kind struct {
	ID        int64  `sql:"id"         json:"id"         form:"id"         query:"id"`
	Name      string `sql:"name"       json:"name"       form:"name"       query:"name"`
	ShortName string `sql:"short_name" json:"short_name" form:"short_name" query:"short_name"`
	Note      string `sql:"note,null"  json:"note"       form:"note"       query:"note"`
	CreatedAt string `sql:"created_at" json:"-"`
	UpdatedAt string `sql:"updated_at" json:"-"`
}

// KindList - struct for kind list
type KindList struct {
	ID        int64  `sql:"id"         json:"id"         form:"id"         query:"id"`
	Name      string `sql:"name"       json:"name"       form:"name"       query:"name"`
	ShortName string `sql:"short_name" json:"short_name" form:"short_name" query:"short_name"`
	Note      string `sql:"note,null"  json:"note"       form:"note"       query:"note"`
}

// GetKind - get one kind by id
func (e *Edb) GetKind(id int64) (Kind, error) {
	var kind Kind
	if id == 0 {
		return kind, nil
	}
	err := e.db.Model(&kind).
		Where("id = ?", id).
		Select()
	if err != nil {
		errmsg("GetKind select", err)
	}
	return kind, err
}

// GetKindList - get kind for list by id
func (e *Edb) GetKindList(id int64) (KindList, error) {
	var kind KindList
	err := e.db.Model(&Kind{}).
		Column("id", "name", "short_name", "note").
		Where("id = ?", id).
		Select(&kind)
	if err != nil {
		errmsg("GetKindList select", err)
	}
	return kind, err
}

// GetKindListAll - get all kind for list
func (e *Edb) GetKindListAll() ([]KindList, error) {
	var kinds []KindList
	err := e.db.Model(&Kind{}).
		Column("id", "name", "short_name", "note").
		Order("name ASC").
		Select(&kinds)
	if err != nil {
		errmsg("GetKindListAll select", err)
	}
	return kinds, err
}

// GetKindSelectAll - get all kind for select
func (e *Edb) GetKindSelectAll() ([]SelectItem, error) {
	var kinds []SelectItem
	err := e.db.Model(&Kind{}).
		Column("id", "name").
		Order("name ASC").
		Select(&kinds)
	if err != nil {
		errmsg("GetKindSelectAll select", err)
	}
	return kinds, err
}

// CreateKind - create new kind
func (e *Edb) CreateKind(kind Kind) (int64, error) {
	err := e.db.Insert(&kind)
	if err != nil {
		errmsg("CreateKind insert", err)
	}
	return kind.ID, nil
}

// UpdateKind - save kind changes
func (e *Edb) UpdateKind(kind Kind) error {
	err := e.db.Update(&kind)
	if err != nil {
		errmsg("UpdateKind update", err)
	}
	return err
}

// DeleteKind - delete kind by id
func (e *Edb) DeleteKind(id int64) error {
	if id == 0 {
		return nil
	}
	_, err := e.db.Model(&Kind{}).
		Where("id = ?", id).
		Delete()
	if err != nil {
		errmsg("DeleteKind delete", err)
	}
	return err
}

func (e *Edb) kindCreateTable() error {
	str := `
		CREATE TABLE IF NOT EXISTS
			kinds (
				id bigserial primary key,
				name text,
				short_name text,
				note text,
				created_at TIMESTAMP without time zone,
				updated_at TIMESTAMP without time zone default now(),
				UNIQUE(name)
			)
	`
	_, err := e.db.Exec(str)
	if err != nil {
		errmsg("kindCreateTable exec", err)
	}
	return err
}
<|endoftext|>