<commit_before>package webClient

import (
  "net/http"
  "github.com/gorilla/websocket"
  "github.com/satori/go.uuid"
  "github.com/M-O-S-E-S/mgm2/core"
)

type clientResponse struct {
  MessageID int
  MessageType string
  Message interface{}
}

type WebsocketConnector struct {
  httpConnector *HttpConnector
  session chan<- core.UserSession
  logger Logger
}

func NewWebsocketConnector(hc *HttpConnector, session chan<- core.UserSession, logger Logger) (*WebsocketConnector) {
  return &WebsocketConnector{hc, session, logger}
}

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

func (wc WebsocketConnector) WebsocketHandler(w http.ResponseWriter, r *http.Request) {

  session, _ := wc.httpConnector.store.Get(r, "MGM")
  // test if session exists
  if len(session.Values) == 0 {
    wc.logger.Info("Websocket closed, no existing session")
    return
  }
  // test origin, etc for websocket security
  // not sure if necessary, we will be over https, and the session is valid

  ws, err := upgrader.Upgrade(w, r, nil)
  if err != nil {
    wc.logger.Error("Error upgrading websocket: %v", err)
    return
  }

  guid := session.Values["guid"].(uuid.UUID)
  uLevel := session.Values["ulevel"].(uint8)

  c := client{ws, make(chan []byte, 64), make(chan []byte, 64), guid, uLevel, wc.logger}
  go c.reader()
  go c.writer()
  wc.session <- c
}

func (c *client) reader() {
  for {
    _, message, err := c.ws.ReadMessage()
    if err != nil {
      break
    }
    c.fromClient<-message
  }
  close(c.fromClient)
  c.ws.Close()
}

func (c *client) writer() {
  for message := range c.toClient {

    err := c.ws.WriteMessage(websocket.TextMessage, message)
    if err != nil {
      break
    }
  }
  close(c.toClient)
  c.ws.Close()
}
<commit_msg>Added json response to session test for websocket connections.  A sessionless client should pick this up and try to authenticate before connecting again.<commit_after>package webClient

import (
  "net/http"
  "github.com/gorilla/websocket"
  "github.com/satori/go.uuid"
  "github.com/M-O-S-E-S/mgm2/core"
  "encoding/json"
)

type clientResponse struct {
  MessageID int
  MessageType string
  Message interface{}
}

type WebsocketConnector struct {
  httpConnector *HttpConnector
  session chan<- core.UserSession
  logger Logger
}

func NewWebsocketConnector(hc *HttpConnector, session chan<- core.UserSession, logger Logger) (*WebsocketConnector) {
  return &WebsocketConnector{hc, session, logger}
}

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

func (wc WebsocketConnector) WebsocketHandler(w http.ResponseWriter, r *http.Request) {

  // test if session exists
  session, _ := wc.httpConnector.store.Get(r, "MGM")
  if len(session.Values) == 0 {
    wc.logger.Info("Websocket closed, no existing session")

    response := clientResponse{ MessageType: "AccessDenied", Message: "No Session Found"}
    js, err := json.Marshal(response)
    if err != nil {
      http.Error(w, err.Error(), http.StatusInternalServerError)
      return
    }
    w.Header().Set("Content-Type", "application/json")
    w.Write(js)
    return
  }
  // test origin, etc for websocket security
  // not sure if necessary, we will be over https, and the session is valid

  ws, err := upgrader.Upgrade(w, r, nil)
  if err != nil {
    wc.logger.Error("Error upgrading websocket: %v", err)
    return
  }

  guid := session.Values["guid"].(uuid.UUID)
  uLevel := session.Values["ulevel"].(uint8)

  c := client{ws, make(chan []byte, 64), make(chan []byte, 64), guid, uLevel, wc.logger}
  go c.reader()
  go c.writer()
  wc.session <- c
}

func (c *client) reader() {
  for {
    _, message, err := c.ws.ReadMessage()
    if err != nil {
      break
    }
    c.fromClient<-message
  }
  close(c.fromClient)
  c.ws.Close()
}

func (c *client) writer() {
  for message := range c.toClient {

    err := c.ws.WriteMessage(websocket.TextMessage, message)
    if err != nil {
      break
    }
  }
  close(c.toClient)
  c.ws.Close()
}
<|endoftext|>


<commit_before>// Copyright 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 routing

import (
	"net/http"

	"github.com/matrix-org/dendrite/clientapi/auth"
	"github.com/matrix-org/dendrite/clientapi/auth/authtypes"

	"github.com/gorilla/mux"
	"github.com/matrix-org/dendrite/clientapi/auth/storage/devices"
	"github.com/matrix-org/dendrite/common"
	"github.com/matrix-org/dendrite/common/config"
	"github.com/matrix-org/dendrite/mediaapi/storage"
	"github.com/matrix-org/dendrite/mediaapi/types"
	"github.com/matrix-org/gomatrixserverlib"
	"github.com/matrix-org/util"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

const pathPrefixR0 = "/_matrix/media/r0"

// Setup registers the media API HTTP handlers
//
// Due to Setup being used to call many other functions, a gocyclo nolint is
// applied:
// nolint: gocyclo
func Setup(
	apiMux *mux.Router,
	cfg *config.Dendrite,
	db *storage.Database,
	deviceDB *devices.Database,
	client *gomatrixserverlib.Client,
) {
	r0mux := apiMux.PathPrefix(pathPrefixR0).Subrouter()

	activeThumbnailGeneration := &types.ActiveThumbnailGeneration{
		PathToResult: map[string]*types.ThumbnailGenerationResult{},
	}
	authData := auth.Data{
		AccountDB:   nil,
		DeviceDB:    deviceDB,
		AppServices: nil,
	}

	// TODO: Add AS support
	r0mux.Handle("/upload", common.MakeAuthAPI(
		"upload", authData,
		func(req *http.Request, _ *authtypes.Device) util.JSONResponse {
			return Upload(req, cfg, db, activeThumbnailGeneration)
		},
	)).Methods(http.MethodPost, http.MethodOptions)

	activeRemoteRequests := &types.ActiveRemoteRequests{
		MXCToResult: map[string]*types.RemoteRequestResult{},
	}
	r0mux.Handle("/download/{serverName}/{mediaId}",
		makeDownloadAPI("download", cfg, db, client, activeRemoteRequests, activeThumbnailGeneration),
	).Methods(http.MethodGet, http.MethodOptions)
	r0mux.Handle("/thumbnail/{serverName}/{mediaId}",
		makeDownloadAPI("thumbnail", cfg, db, client, activeRemoteRequests, activeThumbnailGeneration),
	).Methods(http.MethodGet, http.MethodOptions)
}

func makeDownloadAPI(
	name string,
	cfg *config.Dendrite,
	db *storage.Database,
	client *gomatrixserverlib.Client,
	activeRemoteRequests *types.ActiveRemoteRequests,
	activeThumbnailGeneration *types.ActiveThumbnailGeneration,
) http.HandlerFunc {
	return promhttp.InstrumentHandlerCounter(
		promauto.NewCounterVec(
			prometheus.CounterOpts{
				Name: name,
				Help: "Total number of media_api requests for either thumbnails or full downloads",
			},
			[]string{"code"},
		), http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
			req = util.RequestWithLogging(req)

			// Set common headers returned regardless of the outcome of the request
			util.SetCORSHeaders(w)
			// Content-Type will be overridden in case of returning file data, else we respond with JSON-formatted errors
			w.Header().Set("Content-Type", "application/json")
			vars := mux.Vars(req)
			Download(
				w,
				req,
				gomatrixserverlib.ServerName(vars["serverName"]),
				types.MediaID(vars["mediaId"]),
				cfg,
				db,
				client,
				activeRemoteRequests,
				activeThumbnailGeneration,
				name == "thumbnail",
			)
		},
		))
}
<commit_msg>Refactor InstrumentHandlerCounter definition<commit_after>// Copyright 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 routing

import (
	"net/http"

	"github.com/matrix-org/dendrite/clientapi/auth"
	"github.com/matrix-org/dendrite/clientapi/auth/authtypes"

	"github.com/gorilla/mux"
	"github.com/matrix-org/dendrite/clientapi/auth/storage/devices"
	"github.com/matrix-org/dendrite/common"
	"github.com/matrix-org/dendrite/common/config"
	"github.com/matrix-org/dendrite/mediaapi/storage"
	"github.com/matrix-org/dendrite/mediaapi/types"
	"github.com/matrix-org/gomatrixserverlib"
	"github.com/matrix-org/util"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

const pathPrefixR0 = "/_matrix/media/r0"

// Setup registers the media API HTTP handlers
//
// Due to Setup being used to call many other functions, a gocyclo nolint is
// applied:
// nolint: gocyclo
func Setup(
	apiMux *mux.Router,
	cfg *config.Dendrite,
	db *storage.Database,
	deviceDB *devices.Database,
	client *gomatrixserverlib.Client,
) {
	r0mux := apiMux.PathPrefix(pathPrefixR0).Subrouter()

	activeThumbnailGeneration := &types.ActiveThumbnailGeneration{
		PathToResult: map[string]*types.ThumbnailGenerationResult{},
	}
	authData := auth.Data{
		AccountDB:   nil,
		DeviceDB:    deviceDB,
		AppServices: nil,
	}

	// TODO: Add AS support
	r0mux.Handle("/upload", common.MakeAuthAPI(
		"upload", authData,
		func(req *http.Request, _ *authtypes.Device) util.JSONResponse {
			return Upload(req, cfg, db, activeThumbnailGeneration)
		},
	)).Methods(http.MethodPost, http.MethodOptions)

	activeRemoteRequests := &types.ActiveRemoteRequests{
		MXCToResult: map[string]*types.RemoteRequestResult{},
	}
	r0mux.Handle("/download/{serverName}/{mediaId}",
		makeDownloadAPI("download", cfg, db, client, activeRemoteRequests, activeThumbnailGeneration),
	).Methods(http.MethodGet, http.MethodOptions)
	r0mux.Handle("/thumbnail/{serverName}/{mediaId}",
		makeDownloadAPI("thumbnail", cfg, db, client, activeRemoteRequests, activeThumbnailGeneration),
	).Methods(http.MethodGet, http.MethodOptions)
}

func makeDownloadAPI(
	name string,
	cfg *config.Dendrite,
	db *storage.Database,
	client *gomatrixserverlib.Client,
	activeRemoteRequests *types.ActiveRemoteRequests,
	activeThumbnailGeneration *types.ActiveThumbnailGeneration,
) http.HandlerFunc {
	counterVec := promauto.NewCounterVec(
		prometheus.CounterOpts{
			Name: name,
			Help: "Total number of media_api requests for either thumbnails or full downloads",
		},
		[]string{"code"},
	)
	httpHandler := func(w http.ResponseWriter, req *http.Request) {
		req = util.RequestWithLogging(req)

		// Set common headers returned regardless of the outcome of the request
		util.SetCORSHeaders(w)
		// Content-Type will be overridden in case of returning file data, else we respond with JSON-formatted errors
		w.Header().Set("Content-Type", "application/json")
		vars := mux.Vars(req)
		Download(
			w,
			req,
			gomatrixserverlib.ServerName(vars["serverName"]),
			types.MediaID(vars["mediaId"]),
			cfg,
			db,
			client,
			activeRemoteRequests,
			activeThumbnailGeneration,
			name == "thumbnail",
		)
	}
	return promhttp.InstrumentHandlerCounter(counterVec, http.HandlerFunc(httpHandler))
}
<|endoftext|>


<commit_before>package memaccess

import (
	"fmt"
	"github.com/mozilla/masche/test"
	"os"
	"testing"
)

func TestNewProcessMemoryReader(t *testing.T) {
	cmd, err := test.LaunchTestCase()
	if err != nil {
		t.Fatal(err)
	}
	defer cmd.Process.Kill()

	pid := uint(cmd.Process.Pid)
	reader, err, softerrors := NewProcessMemoryReader(pid)
	test.PrintSoftErrors(softerrors)
	if err != nil {
		t.Fatal(err)
	}
	defer reader.Close()
}

func TestManuallyWalk(t *testing.T) {
	cmd, err := test.LaunchTestCase()
	if err != nil {
		t.Fatal(err)
	}
	defer cmd.Process.Kill()

	pid := uint(cmd.Process.Pid)
	reader, err, softerrors := NewProcessMemoryReader(pid)
	test.PrintSoftErrors(softerrors)
	if err != nil {
		t.Fatal(err)
	}
	defer reader.Close()

	var region MemoryRegion
	region, err, softerrors = reader.NextReadableMemoryRegion(0)
	test.PrintSoftErrors(softerrors)
	if err != nil {
		t.Fatal(err)
	}

	if region == NoRegionAvailable {
		t.Error("No starting region returned")
	}

	previousRegion := region
	for region != NoRegionAvailable {
		region, err, softerrors = reader.NextReadableMemoryRegion(region.Address + uintptr(region.Size))
		test.PrintSoftErrors(softerrors)
		if err != nil {
			t.Fatal(err)
		}

		if region != NoRegionAvailable && region.Address < previousRegion.Address+uintptr(previousRegion.Size) {
			t.Error("Returned region is not after the previous one.")
		}

		previousRegion = region
	}
}

func TestCopyMemory(t *testing.T) {
	cmd, err := test.LaunchTestCaseAndWaitForInitialization()
	if err != nil {
		t.Fatal(err)
	}
	defer cmd.Process.Kill()

	pid := uint(cmd.Process.Pid)
	reader, err, softerrors := NewProcessMemoryReader(pid)
	test.PrintSoftErrors(softerrors)
	if err != nil {
		t.Fatal(err)
	}
	defer reader.Close()

	var region MemoryRegion
	region, err, softerrors = reader.NextReadableMemoryRegion(0)
	test.PrintSoftErrors(softerrors)
	if err != nil {
		t.Fatal(err)
	}

	if region == NoRegionAvailable {
		t.Error("No starting region returned")
	}

	min_region_size := uint(os.Getpagesize() + 100) // one page plus something

	for region.Size < min_region_size {
		if region == NoRegionAvailable {
			t.Fatal("We couldn't find a region of %d bytes", min_region_size)
		}

		region, err, softerrors = reader.NextReadableMemoryRegion(region.Address + uintptr(region.Size))
		test.PrintSoftErrors(softerrors)
		if err != nil {
			t.Fatal(err)
		}
	}

	buffers := [][]byte{
		make([]byte, 2),
		make([]byte, os.Getpagesize()),
		make([]byte, min_region_size),
	}

	for _, buffer := range buffers {
		// Valid read
		err, softerrors = reader.CopyMemory(region.Address, buffer)
		test.PrintSoftErrors(softerrors)
		if err != nil {
			t.Error(fmt.Sprintf("Couldn't read %d bytes from region", len(buffer)))
		}

		// Crossing boundaries
		err, softerrors = reader.CopyMemory(region.Address+uintptr(region.Size)-uintptr(len(buffer)/2), buffer)
		test.PrintSoftErrors(softerrors)
		if err == nil {
			t.Error(fmt.Sprintf("Read %d bytes inbetween regions", len(buffer)))
		}

		// Entirely outside region
		err, softerrors = reader.CopyMemory(region.Address+uintptr(region.Size), buffer)
		test.PrintSoftErrors(softerrors)
		if err == nil {
			t.Error(fmt.Sprintf("Read %d bytes after the region", len(buffer)))
		}
	}

}

func memoryRegionsOverlap(region1 MemoryRegion, region2 MemoryRegion) bool {
	region1End := region1.Address + uintptr(region1.Size)
	region2End := region2.Address + uintptr(region2.Size)

	if region2.Address >= region1.Address {
		if region2.Address < region1End {
			return true
		}
	} else {
		if region2End <= region1End {
			return true
		}
	}

	return false
}

func TestWalkMemoryDoesntOverlapTheBuffer(t *testing.T) {
	cmd, err := test.LaunchTestCaseAndWaitForInitialization()
	if err != nil {
		t.Fatal(err)
	}
	defer cmd.Process.Kill()

	pid := uint(cmd.Process.Pid)
	reader, err, softerrors := NewProcessMemoryReader(pid)
	test.PrintSoftErrors(softerrors)
	if err != nil {
		t.Fatal(err)
	}
	defer reader.Close()

	pageSize := uint(os.Getpagesize())
	bufferSizes := []uint{1024, pageSize, pageSize + 100, pageSize * 2, pageSize*2 + 123}
	for _, size := range bufferSizes {

		lastRegion := MemoryRegion{}
		err, softerrors = WalkMemory(reader, 0, size, func(address uintptr, buffer []byte) (keepSearching bool) {
			currenRegion := MemoryRegion{Address: address, Size: uint(len(buffer))}
			if memoryRegionsOverlap(lastRegion, currenRegion) {
				t.Errorf("Regions overlap while reading %d at a time: %v %v", size, lastRegion, currenRegion)
				return false
			}

			lastRegion = currenRegion
			return true
		})
		test.PrintSoftErrors(softerrors)
		if err != nil {
			t.Fatal(err)
		}
	}
}

func TestWalkRegionReadsEntireRegion(t *testing.T) {
	cmd, err := test.LaunchTestCaseAndWaitForInitialization()
	if err != nil {
		t.Fatal(err)
	}
	defer cmd.Process.Kill()

	pid := uint(cmd.Process.Pid)
	reader, err, softerrors := NewProcessMemoryReader(pid)
	test.PrintSoftErrors(softerrors)
	if err != nil {
		t.Fatal(err)
	}
	defer reader.Close()

	pageSize := uint(os.Getpagesize())
	bufferSizes := []uint{1024, pageSize, pageSize + 100, pageSize * 2, pageSize*2 + 123}

	var region MemoryRegion
	region, err, softerrors = reader.NextReadableMemoryRegion(0)
	test.PrintSoftErrors(softerrors)
	if err != nil {
		t.Fatal(err)
	}

	if region == NoRegionAvailable {
		t.Error("No starting region returned")
	}

	min_region_size := bufferSizes[len(bufferSizes)-1]
	for region.Size < min_region_size {
		if region == NoRegionAvailable {
			t.Fatal("We couldn't find a region of %d bytes", min_region_size)
		}

		region, err, softerrors = reader.NextReadableMemoryRegion(region.Address + uintptr(region.Size))
		test.PrintSoftErrors(softerrors)
		if err != nil {
			t.Fatal(err)
		}
	}

	for _, size := range bufferSizes {
		buf := make([]byte, size)
		readRegion := MemoryRegion{}

		_, _, err, softerrors := walkRegion(reader, region, buf,
			func(address uintptr, buffer []byte) (keepSearching bool) {
				if readRegion.Address == 0 {
					readRegion.Address = address
					readRegion.Size = uint(len(buffer))
					return true
				}

				readRegionLimit := readRegion.Address + uintptr(readRegion.Size)
				if readRegionLimit != address {
					t.Error(fmt.Sprintf("walkRegion skept %d bytes starting at %x", address-readRegionLimit,
						readRegionLimit))
					return false
				}

				readRegion.Size += uint(len(buffer))
				return true
			})
		test.PrintSoftErrors(softerrors)
		if err != nil {
			t.Fatal(err)
		}
	}
}
<commit_msg>[memaccess] Fix typo<commit_after>package memaccess

import (
	"fmt"
	"github.com/mozilla/masche/test"
	"os"
	"testing"
)

func TestNewProcessMemoryReader(t *testing.T) {
	cmd, err := test.LaunchTestCase()
	if err != nil {
		t.Fatal(err)
	}
	defer cmd.Process.Kill()

	pid := uint(cmd.Process.Pid)
	reader, err, softerrors := NewProcessMemoryReader(pid)
	test.PrintSoftErrors(softerrors)
	if err != nil {
		t.Fatal(err)
	}
	defer reader.Close()
}

func TestManuallyWalk(t *testing.T) {
	cmd, err := test.LaunchTestCase()
	if err != nil {
		t.Fatal(err)
	}
	defer cmd.Process.Kill()

	pid := uint(cmd.Process.Pid)
	reader, err, softerrors := NewProcessMemoryReader(pid)
	test.PrintSoftErrors(softerrors)
	if err != nil {
		t.Fatal(err)
	}
	defer reader.Close()

	var region MemoryRegion
	region, err, softerrors = reader.NextReadableMemoryRegion(0)
	test.PrintSoftErrors(softerrors)
	if err != nil {
		t.Fatal(err)
	}

	if region == NoRegionAvailable {
		t.Error("No starting region returned")
	}

	previousRegion := region
	for region != NoRegionAvailable {
		region, err, softerrors = reader.NextReadableMemoryRegion(region.Address + uintptr(region.Size))
		test.PrintSoftErrors(softerrors)
		if err != nil {
			t.Fatal(err)
		}

		if region != NoRegionAvailable && region.Address < previousRegion.Address+uintptr(previousRegion.Size) {
			t.Error("Returned region is not after the previous one.")
		}

		previousRegion = region
	}
}

func TestCopyMemory(t *testing.T) {
	cmd, err := test.LaunchTestCaseAndWaitForInitialization()
	if err != nil {
		t.Fatal(err)
	}
	defer cmd.Process.Kill()

	pid := uint(cmd.Process.Pid)
	reader, err, softerrors := NewProcessMemoryReader(pid)
	test.PrintSoftErrors(softerrors)
	if err != nil {
		t.Fatal(err)
	}
	defer reader.Close()

	var region MemoryRegion
	region, err, softerrors = reader.NextReadableMemoryRegion(0)
	test.PrintSoftErrors(softerrors)
	if err != nil {
		t.Fatal(err)
	}

	if region == NoRegionAvailable {
		t.Error("No starting region returned")
	}

	min_region_size := uint(os.Getpagesize() + 100) // one page plus something

	for region.Size < min_region_size {
		if region == NoRegionAvailable {
			t.Fatal("We couldn't find a region of %d bytes", min_region_size)
		}

		region, err, softerrors = reader.NextReadableMemoryRegion(region.Address + uintptr(region.Size))
		test.PrintSoftErrors(softerrors)
		if err != nil {
			t.Fatal(err)
		}
	}

	buffers := [][]byte{
		make([]byte, 2),
		make([]byte, os.Getpagesize()),
		make([]byte, min_region_size),
	}

	for _, buffer := range buffers {
		// Valid read
		err, softerrors = reader.CopyMemory(region.Address, buffer)
		test.PrintSoftErrors(softerrors)
		if err != nil {
			t.Error(fmt.Sprintf("Couldn't read %d bytes from region", len(buffer)))
		}

		// Crossing boundaries
		err, softerrors = reader.CopyMemory(region.Address+uintptr(region.Size)-uintptr(len(buffer)/2), buffer)
		test.PrintSoftErrors(softerrors)
		if err == nil {
			t.Error(fmt.Sprintf("Read %d bytes inbetween regions", len(buffer)))
		}

		// Entirely outside region
		err, softerrors = reader.CopyMemory(region.Address+uintptr(region.Size), buffer)
		test.PrintSoftErrors(softerrors)
		if err == nil {
			t.Error(fmt.Sprintf("Read %d bytes after the region", len(buffer)))
		}
	}

}

func memoryRegionsOverlap(region1 MemoryRegion, region2 MemoryRegion) bool {
	region1End := region1.Address + uintptr(region1.Size)
	region2End := region2.Address + uintptr(region2.Size)

	if region2.Address >= region1.Address {
		if region2.Address < region1End {
			return true
		}
	} else {
		if region2End <= region1End {
			return true
		}
	}

	return false
}

func TestWalkMemoryDoesntOverlapTheBuffer(t *testing.T) {
	cmd, err := test.LaunchTestCaseAndWaitForInitialization()
	if err != nil {
		t.Fatal(err)
	}
	defer cmd.Process.Kill()

	pid := uint(cmd.Process.Pid)
	reader, err, softerrors := NewProcessMemoryReader(pid)
	test.PrintSoftErrors(softerrors)
	if err != nil {
		t.Fatal(err)
	}
	defer reader.Close()

	pageSize := uint(os.Getpagesize())
	bufferSizes := []uint{1024, pageSize, pageSize + 100, pageSize * 2, pageSize*2 + 123}
	for _, size := range bufferSizes {

		lastRegion := MemoryRegion{}
		err, softerrors = WalkMemory(reader, 0, size, func(address uintptr, buffer []byte) (keepSearching bool) {
			currentRegion := MemoryRegion{Address: address, Size: uint(len(buffer))}
			if memoryRegionsOverlap(lastRegion, currentRegion) {
				t.Errorf("Regions overlap while reading %d at a time: %v %v", size, lastRegion, currentRegion)
				return false
			}

			lastRegion = currentRegion
			return true
		})
		test.PrintSoftErrors(softerrors)
		if err != nil {
			t.Fatal(err)
		}
	}
}

func TestWalkRegionReadsEntireRegion(t *testing.T) {
	cmd, err := test.LaunchTestCaseAndWaitForInitialization()
	if err != nil {
		t.Fatal(err)
	}
	defer cmd.Process.Kill()

	pid := uint(cmd.Process.Pid)
	reader, err, softerrors := NewProcessMemoryReader(pid)
	test.PrintSoftErrors(softerrors)
	if err != nil {
		t.Fatal(err)
	}
	defer reader.Close()

	pageSize := uint(os.Getpagesize())
	bufferSizes := []uint{1024, pageSize, pageSize + 100, pageSize * 2, pageSize*2 + 123}

	var region MemoryRegion
	region, err, softerrors = reader.NextReadableMemoryRegion(0)
	test.PrintSoftErrors(softerrors)
	if err != nil {
		t.Fatal(err)
	}

	if region == NoRegionAvailable {
		t.Error("No starting region returned")
	}

	min_region_size := bufferSizes[len(bufferSizes)-1]
	for region.Size < min_region_size {
		if region == NoRegionAvailable {
			t.Fatal("We couldn't find a region of %d bytes", min_region_size)
		}

		region, err, softerrors = reader.NextReadableMemoryRegion(region.Address + uintptr(region.Size))
		test.PrintSoftErrors(softerrors)
		if err != nil {
			t.Fatal(err)
		}
	}

	for _, size := range bufferSizes {
		buf := make([]byte, size)
		readRegion := MemoryRegion{}

		_, _, err, softerrors := walkRegion(reader, region, buf,
			func(address uintptr, buffer []byte) (keepSearching bool) {
				if readRegion.Address == 0 {
					readRegion.Address = address
					readRegion.Size = uint(len(buffer))
					return true
				}

				readRegionLimit := readRegion.Address + uintptr(readRegion.Size)
				if readRegionLimit != address {
					t.Error(fmt.Sprintf("walkRegion skept %d bytes starting at %x", address-readRegionLimit,
						readRegionLimit))
					return false
				}

				readRegion.Size += uint(len(buffer))
				return true
			})
		test.PrintSoftErrors(softerrors)
		if err != nil {
			t.Fatal(err)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. 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 sources

import (
	"fmt"
	"sync"
	"time"

	"github.com/golang/glog"
	"k8s.io/heapster/sources/api"
	"k8s.io/heapster/sources/datasource"
	"k8s.io/heapster/sources/nodes"
)

type kubeNodeMetrics struct {
	kubeletApi  datasource.Kubelet
	kubeletPort int
	nodesApi    nodes.NodesApi
}

func NewKubeNodeMetrics(kubeletPort int, kubeletApi datasource.Kubelet, nodesApi nodes.NodesApi) api.Source {
	return &kubeNodeMetrics{
		kubeletApi:  kubeletApi,
		kubeletPort: kubeletPort,
		nodesApi:    nodesApi,
	}
}

const (
	rootContainer             = "/"
	KubeNodeMetricsSourceName = "Kube Node Metrics Source"
)

var knownContainers = map[string]string{
	"/docker-daemon": "docker-daemon",
	"/kubelet":       "kubelet",
	"/kube-proxy":    "kube-proxy",
	"/system":        "system",
}

// Returns the host container, non-Kubernetes containers, and an error (if any).
func (self *kubeNodeMetrics) updateStats(host nodes.Host, info nodes.Info, start, end time.Time) (*api.Container, []api.Container, error) {
	// Get information for all containers.
	containers, err := self.kubeletApi.GetAllRawContainers(datasource.Host{IP: info.InternalIP, Port: self.kubeletPort}, start, end)
	if err != nil {
		glog.V(3).Infof("Failed to get container stats from Kubelet on node %q", host)
		return nil, []api.Container{}, fmt.Errorf("failed to get container stats from Kubelet on node %q: %v", host, err)
	}
	if len(containers) == 0 {
		// no stats found.
		glog.V(3).Infof("No container stats from Kubelet on node %q", host)
		return nil, []api.Container{}, fmt.Errorf("no container stats from Kubelet on node %q", host)
	}

	// Find host container.
	hostIndex := -1
	hostString := string(host)
	externalID := string(info.ExternalID)
	for i := range containers {
		if containers[i].Name == rootContainer {
			hostIndex = i
		}
		if newName, exists := knownContainers[containers[i].Name]; exists {
			containers[i].Name = newName
		}
		containers[i].Hostname = hostString
		containers[i].ExternalID = externalID
	}
	var hostContainer *api.Container
	if hostIndex >= 0 {
		hostCopy := containers[hostIndex]
		hostContainer = &hostCopy
		containers = append(containers[:hostIndex], containers[hostIndex+1:]...)
	}
	// This is temporary workaround for #399. To make unit consistent with cadvisor normalize to a conversion factor of 1024.
	hostContainer.Spec.Cpu.Limit = info.CpuCapacity * 1024 / 1000
	hostContainer.Spec.Memory.Limit = info.MemCapacity
	return hostContainer, containers, nil
}

// Returns the host containers, non-Kubernetes containers, and an error (if any).
func (self *kubeNodeMetrics) getNodesInfo(nodeList *nodes.NodeList, start, end time.Time) ([]api.Container, []api.Container, error) {
	var (
		lock sync.Mutex
		wg   sync.WaitGroup
	)
	hostContainers := make([]api.Container, 0, len(nodeList.Items))
	rawContainers := make([]api.Container, 0, len(nodeList.Items))
	for host, info := range nodeList.Items {
		wg.Add(1)
		go func(host nodes.Host, info nodes.Info) {
			defer wg.Done()
			if hostContainer, containers, err := self.updateStats(host, info, start, end); err == nil {
				lock.Lock()
				defer lock.Unlock()
				if hostContainers != nil {
					hostContainers = append(hostContainers, *hostContainer)
				}
				rawContainers = append(rawContainers, containers...)
			}
		}(host, info)
	}
	wg.Wait()

	return hostContainers, rawContainers, nil
}

func (self *kubeNodeMetrics) GetInfo(start, end time.Time) (api.AggregateData, error) {
	kubeNodes, err := self.nodesApi.List()
	if err != nil || len(kubeNodes.Items) == 0 {
		return api.AggregateData{}, err
	}
	glog.V(3).Info("Fetched list of nodes from the master")
	hostContainers, rawContainers, err := self.getNodesInfo(kubeNodes, start, end)
	if err != nil {
		return api.AggregateData{}, err
	}

	return api.AggregateData{
		Machine:    hostContainers,
		Containers: rawContainers,
	}, nil
}

func (self *kubeNodeMetrics) DebugInfo() string {
	desc := "Source type: Kube Node Metrics\n"
	desc += self.nodesApi.DebugInfo() + "\n"

	return desc
}

func (kns *kubeNodeMetrics) Name() string {
	return KubeNodeMetricsSourceName
}
<commit_msg>Fix nil pointer reference panic in kube_nodes<commit_after>// Copyright 2015 Google Inc. 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 sources

import (
	"fmt"
	"sync"
	"time"

	"github.com/golang/glog"
	"k8s.io/heapster/sources/api"
	"k8s.io/heapster/sources/datasource"
	"k8s.io/heapster/sources/nodes"
)

type kubeNodeMetrics struct {
	kubeletApi  datasource.Kubelet
	kubeletPort int
	nodesApi    nodes.NodesApi
}

func NewKubeNodeMetrics(kubeletPort int, kubeletApi datasource.Kubelet, nodesApi nodes.NodesApi) api.Source {
	return &kubeNodeMetrics{
		kubeletApi:  kubeletApi,
		kubeletPort: kubeletPort,
		nodesApi:    nodesApi,
	}
}

const (
	rootContainer             = "/"
	KubeNodeMetricsSourceName = "Kube Node Metrics Source"
)

var knownContainers = map[string]string{
	"/docker-daemon": "docker-daemon",
	"/kubelet":       "kubelet",
	"/kube-proxy":    "kube-proxy",
	"/system":        "system",
}

// Returns the host container, non-Kubernetes containers, and an error (if any).
func (self *kubeNodeMetrics) updateStats(host nodes.Host, info nodes.Info, start, end time.Time) (*api.Container, []api.Container, error) {
	// Get information for all containers.
	containers, err := self.kubeletApi.GetAllRawContainers(datasource.Host{IP: info.InternalIP, Port: self.kubeletPort}, start, end)
	if err != nil {
		glog.V(3).Infof("Failed to get container stats from Kubelet on node %q", host)
		return nil, []api.Container{}, fmt.Errorf("failed to get container stats from Kubelet on node %q: %v", host, err)
	}
	if len(containers) == 0 {
		// no stats found.
		glog.V(3).Infof("No container stats from Kubelet on node %q", host)
		return nil, []api.Container{}, fmt.Errorf("no container stats from Kubelet on node %q", host)
	}

	// Find host container.
	hostIndex := -1
	hostString := string(host)
	externalID := string(info.ExternalID)
	for i := range containers {
		if containers[i].Name == rootContainer {
			hostIndex = i
		}
		if newName, exists := knownContainers[containers[i].Name]; exists {
			containers[i].Name = newName
		}
		containers[i].Hostname = hostString
		containers[i].ExternalID = externalID
	}
	var hostContainer *api.Container
	if hostIndex >= 0 {
		hostCopy := containers[hostIndex]
		hostContainer = &hostCopy
		containers = append(containers[:hostIndex], containers[hostIndex+1:]...)
		// This is temporary workaround for #399. To make unit consistent with cadvisor normalize to a conversion factor of 1024.
		hostContainer.Spec.Cpu.Limit = info.CpuCapacity * 1024 / 1000
		hostContainer.Spec.Memory.Limit = info.MemCapacity
		return hostContainer, containers, nil
	} else {
		return nil, []api.Container{}, fmt.Errorf("Host container not found")
	}
}

// Returns the host containers, non-Kubernetes containers, and an error (if any).
func (self *kubeNodeMetrics) getNodesInfo(nodeList *nodes.NodeList, start, end time.Time) ([]api.Container, []api.Container, error) {
	var (
		lock sync.Mutex
		wg   sync.WaitGroup
	)
	hostContainers := make([]api.Container, 0, len(nodeList.Items))
	rawContainers := make([]api.Container, 0, len(nodeList.Items))
	for host, info := range nodeList.Items {
		wg.Add(1)
		go func(host nodes.Host, info nodes.Info) {
			defer wg.Done()
			if hostContainer, containers, err := self.updateStats(host, info, start, end); err == nil {
				lock.Lock()
				defer lock.Unlock()
				if hostContainers != nil {
					hostContainers = append(hostContainers, *hostContainer)
				}
				rawContainers = append(rawContainers, containers...)
			}
		}(host, info)
	}
	wg.Wait()

	return hostContainers, rawContainers, nil
}

func (self *kubeNodeMetrics) GetInfo(start, end time.Time) (api.AggregateData, error) {
	kubeNodes, err := self.nodesApi.List()
	if err != nil || len(kubeNodes.Items) == 0 {
		return api.AggregateData{}, err
	}
	glog.V(3).Info("Fetched list of nodes from the master")
	hostContainers, rawContainers, err := self.getNodesInfo(kubeNodes, start, end)
	if err != nil {
		return api.AggregateData{}, err
	}

	return api.AggregateData{
		Machine:    hostContainers,
		Containers: rawContainers,
	}, nil
}

func (self *kubeNodeMetrics) DebugInfo() string {
	desc := "Source type: Kube Node Metrics\n"
	desc += self.nodesApi.DebugInfo() + "\n"

	return desc
}

func (kns *kubeNodeMetrics) Name() string {
	return KubeNodeMetricsSourceName
}
<|endoftext|>


<commit_before>// Copyright 2014 Google Inc. 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 nodes

import (
	"fmt"
	"net"
	"sync"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client/cache"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/fields"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	"github.com/golang/glog"
)

type kubeNodes struct {
	client *client.Client
	// a means to list all minions
	nodeLister *cache.StoreToNodeLister
	reflector  *cache.Reflector
	// Used to stop the existing reflector.
	stopChan   chan struct{}
	goodNodes  []string       // guarded by stateLock
	nodeErrors map[string]int // guarded by stateLock
	stateLock  sync.RWMutex
}

func (self *kubeNodes) recordNodeError(name string) {
	self.stateLock.Lock()
	defer self.stateLock.Unlock()

	self.nodeErrors[name]++
}

func (self *kubeNodes) recordGoodNodes(nodes []string) {
	self.stateLock.Lock()
	defer self.stateLock.Unlock()

	self.goodNodes = nodes
}

func parseSelectorOrDie(s string) labels.Selector {
	selector, err := labels.Parse(s)
	if err != nil {
		panic(err)
	}
	return selector
}

func (self *kubeNodes) getNodeInfoAndHostname(node api.Node) (Info, string) {
	nodeInfo := Info{}
	hostname := ""
	for _, addr := range node.Status.Addresses {
		switch addr.Type {
		case api.NodeExternalIP:
			nodeInfo.PublicIP = addr.Address
		case api.NodeInternalIP:
			nodeInfo.InternalIP = addr.Address
		case api.NodeHostName:
			hostname = addr.Address
		}
	}
	if hostname == "" {
		hostname = node.Name
	}
	if nodeInfo.InternalIP == "" {
		addrs, err := net.LookupIP(hostname)
		if err == nil {
			nodeInfo.InternalIP = addrs[0].String()
		} else {
			glog.Errorf("Skipping host %s since looking up its IP failed - %s", node.Name, err)
			self.recordNodeError(node.Name)
		}
	}
	return nodeInfo, hostname
}

func (self *kubeNodes) List() (*NodeList, error) {
	nodeList := newNodeList()
	allNodes, err := self.nodeLister.List()
	if err != nil {
		glog.Errorf("failed to list minions via watch interface - %v", err)
		return nil, fmt.Errorf("failed to list minions via watch interface - %v", err)
	}
	glog.V(5).Infof("all kube nodes: %+v", allNodes)

	goodNodes := []string{}
	for _, node := range allNodes.Items {
		nodeInfo, hostname := self.getNodeInfoAndHostname(node)

		nodeList.Items[Host(hostname)] = nodeInfo
		goodNodes = append(goodNodes, node.Name)
	}
	self.recordGoodNodes(goodNodes)
	glog.V(5).Infof("kube nodes found: %+v", nodeList)
	return nodeList, nil
}

func (self *kubeNodes) getState() string {
	self.stateLock.RLock()
	defer self.stateLock.RUnlock()

	state := "\tHealthy Nodes:\n"
	for _, node := range self.goodNodes {
		state += fmt.Sprintf("\t\t%s\n", node)
	}
	if len(self.nodeErrors) > 0 {
		state += fmt.Sprintf("\tNode Errors: %+v\n", self.nodeErrors)
	} else {
		state += "\tNo node errors\n"
	}
	return state
}

func (self *kubeNodes) DebugInfo() string {
	desc := "Kubernetes Nodes plugin: \n"
	desc += self.getState()
	desc += "\n"

	return desc
}

func NewKubeNodes(client *client.Client) (NodesApi, error) {
	if client == nil {
		return nil, fmt.Errorf("client is nil")
	}

	lw := cache.NewListWatchFromClient(client, "nodes", api.NamespaceAll, fields.Everything())
	nodeLister := &cache.StoreToNodeLister{Store: cache.NewStore(cache.MetaNamespaceKeyFunc)}
	reflector := cache.NewReflector(lw, &api.Node{}, nodeLister.Store, 0)
	stopChan := make(chan struct{})
	reflector.RunUntil(stopChan)

	return &kubeNodes{
		client:     client,
		nodeLister: nodeLister,
		reflector:  reflector,
		stopChan:   stopChan,
		nodeErrors: make(map[string]int),
	}, nil
}
<commit_msg>Only add nodes without errors to goodNodes list<commit_after>// Copyright 2014 Google Inc. 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 nodes

import (
	"fmt"
	"net"
	"sync"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client/cache"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/fields"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	"github.com/golang/glog"
)

type kubeNodes struct {
	client *client.Client
	// a means to list all minions
	nodeLister *cache.StoreToNodeLister
	reflector  *cache.Reflector
	// Used to stop the existing reflector.
	stopChan   chan struct{}
	goodNodes  []string       // guarded by stateLock
	nodeErrors map[string]int // guarded by stateLock
	stateLock  sync.RWMutex
}

func (self *kubeNodes) recordNodeError(name string) {
	self.stateLock.Lock()
	defer self.stateLock.Unlock()

	self.nodeErrors[name]++
}

func (self *kubeNodes) recordGoodNodes(nodes []string) {
	self.stateLock.Lock()
	defer self.stateLock.Unlock()

	self.goodNodes = nodes
}

func parseSelectorOrDie(s string) labels.Selector {
	selector, err := labels.Parse(s)
	if err != nil {
		panic(err)
	}
	return selector
}

func (self *kubeNodes) getNodeInfoAndHostname(node api.Node) (Info, string, error) {
	nodeInfo := Info{}
	hostname := ""
	var nodeErr error
	for _, addr := range node.Status.Addresses {
		switch addr.Type {
		case api.NodeExternalIP:
			nodeInfo.PublicIP = addr.Address
		case api.NodeInternalIP:
			nodeInfo.InternalIP = addr.Address
		case api.NodeHostName:
			hostname = addr.Address
		}
	}
	if hostname == "" {
		hostname = node.Name
	}
	if nodeInfo.InternalIP == "" {
		addrs, err := net.LookupIP(hostname)
		if err == nil {
			nodeInfo.InternalIP = addrs[0].String()
		} else {
			glog.Errorf("Skipping host %s since looking up its IP failed - %s", node.Name, err)
			self.recordNodeError(node.Name)
			nodeErr = err
		}
	}
	return nodeInfo, hostname, nodeErr
}

func (self *kubeNodes) List() (*NodeList, error) {
	nodeList := newNodeList()
	allNodes, err := self.nodeLister.List()
	if err != nil {
		glog.Errorf("failed to list minions via watch interface - %v", err)
		return nil, fmt.Errorf("failed to list minions via watch interface - %v", err)
	}
	glog.V(5).Infof("all kube nodes: %+v", allNodes)

	goodNodes := []string{}
	for _, node := range allNodes.Items {
		nodeInfo, hostname, err := self.getNodeInfoAndHostname(node)

		nodeList.Items[Host(hostname)] = nodeInfo
		if err == nil {
			goodNodes = append(goodNodes, node.Name)
		}
	}
	self.recordGoodNodes(goodNodes)
	glog.V(5).Infof("kube nodes found: %+v", nodeList)
	return nodeList, nil
}

func (self *kubeNodes) getState() string {
	self.stateLock.RLock()
	defer self.stateLock.RUnlock()

	state := "\tHealthy Nodes:\n"
	for _, node := range self.goodNodes {
		state += fmt.Sprintf("\t\t%s\n", node)
	}
	if len(self.nodeErrors) > 0 {
		state += fmt.Sprintf("\tNode Errors: %+v\n", self.nodeErrors)
	} else {
		state += "\tNo node errors\n"
	}
	return state
}

func (self *kubeNodes) DebugInfo() string {
	desc := "Kubernetes Nodes plugin: \n"
	desc += self.getState()
	desc += "\n"

	return desc
}

func NewKubeNodes(client *client.Client) (NodesApi, error) {
	if client == nil {
		return nil, fmt.Errorf("client is nil")
	}

	lw := cache.NewListWatchFromClient(client, "nodes", api.NamespaceAll, fields.Everything())
	nodeLister := &cache.StoreToNodeLister{Store: cache.NewStore(cache.MetaNamespaceKeyFunc)}
	reflector := cache.NewReflector(lw, &api.Node{}, nodeLister.Store, 0)
	stopChan := make(chan struct{})
	reflector.RunUntil(stopChan)

	return &kubeNodes{
		client:     client,
		nodeLister: nodeLister,
		reflector:  reflector,
		stopChan:   stopChan,
		nodeErrors: make(map[string]int),
	}, nil
}
<|endoftext|>


<commit_before>// Copyright 2016 Google Inc. 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 compact

import (
	"bytes"
	"encoding/base64"
	"errors"
	"fmt"
	"math/bits"
	"strings"
	"testing"

	"github.com/google/trillian/merkle"
	"github.com/google/trillian/merkle/rfc6962"
	"github.com/google/trillian/storage"
	"github.com/google/trillian/testonly"
	"github.com/kylelemons/godebug/pretty"
)

// This check ensures that the compact Merkle tree contains the correct set of
// nodes, i.e. the node on level i is present iff i-th bit of tree size is 1.
func checkUnusedNodesInvariant(t *Tree) error {
	size := t.size
	sizeBits := bits.Len64(uint64(size))
	if got, want := len(t.nodes), sizeBits; got != want {
		return fmt.Errorf("nodes mismatch: have %v nodes, want %v", got, want)
	}
	for level := 0; level < sizeBits; level++ {
		if size&1 == 1 {
			if t.nodes[level] == nil {
				return fmt.Errorf("missing node at level %d", level)
			}
		} else if t.nodes[level] != nil {
			return fmt.Errorf("unexpected node at level %d", level)
		}
		size >>= 1
	}
	return nil
}

func TestAddingLeaves(t *testing.T) {
	inputs := testonly.MerkleTreeLeafTestInputs()
	roots := testonly.MerkleTreeLeafTestRootHashes()
	hashes := testonly.CompactMerkleTreeLeafTestNodeHashes()

	// We test the "same" thing 3 different ways this is to ensure than any lazy
	// update strategy being employed by the implementation doesn't affect the
	// api-visible calculation of root & size.
	{
		// First tree, add nodes one-by-one
		tree := NewTree(rfc6962.DefaultHasher)
		if got, want := tree.Size(), int64(0); got != want {
			t.Errorf("Size()=%d, want %d", got, want)
		}
		if got, want := tree.CurrentRoot(), testonly.EmptyMerkleTreeRootHash(); !bytes.Equal(got, want) {
			t.Errorf("CurrentRoot()=%x, want %x", got, want)
		}

		for i := 0; i < 8; i++ {
			tree.AddLeaf(inputs[i], func(int, int64, []byte) error {
				return nil
			})
			if err := checkUnusedNodesInvariant(tree); err != nil {
				t.Fatalf("UnusedNodesInvariant check failed: %v", err)
			}
			if got, want := tree.Size(), int64(i+1); got != want {
				t.Errorf("Size()=%d, want %d", got, want)
			}
			if got, want := tree.CurrentRoot(), roots[i]; !bytes.Equal(got, want) {
				t.Errorf("CurrentRoot()=%v, want %v", got, want)
			}
			if diff := pretty.Compare(tree.Hashes(), hashes[i]); diff != "" {
				t.Errorf("post-Hashes() diff:\n%v", diff)
			}
		}
	}

	{
		// Second tree, add nodes all at once
		tree := NewTree(rfc6962.DefaultHasher)
		for i := 0; i < 8; i++ {
			tree.AddLeaf(inputs[i], func(int, int64, []byte) error {
				return nil
			})
			if err := checkUnusedNodesInvariant(tree); err != nil {
				t.Fatalf("UnusedNodesInvariant check failed: %v", err)
			}
		}
		if got, want := tree.Size(), int64(8); got != want {
			t.Errorf("Size()=%d, want %d", got, want)
		}
		if got, want := tree.CurrentRoot(), roots[7]; !bytes.Equal(got, want) {
			t.Errorf("CurrentRoot()=%v, want %v", got, want)
		}
		if diff := pretty.Compare(tree.Hashes(), hashes[7]); diff != "" {
			t.Errorf("post-Hashes() diff:\n%v", diff)
		}
	}

	{
		// Third tree, add nodes in two chunks
		tree := NewTree(rfc6962.DefaultHasher)
		for i := 0; i < 3; i++ {
			tree.AddLeaf(inputs[i], func(int, int64, []byte) error {
				return nil
			})
			if err := checkUnusedNodesInvariant(tree); err != nil {
				t.Fatalf("UnusedNodesInvariant check failed: %v", err)
			}
		}
		if got, want := tree.Size(), int64(3); got != want {
			t.Errorf("Size()=%d, want %d", got, want)
		}
		if got, want := tree.CurrentRoot(), roots[2]; !bytes.Equal(got, want) {
			t.Errorf("CurrentRoot()=%v, want %v", got, want)
		}
		if diff := pretty.Compare(tree.Hashes(), hashes[2]); diff != "" {
			t.Errorf("post-Hashes() diff:\n%v", diff)
		}

		for i := 3; i < 8; i++ {
			tree.AddLeaf(inputs[i], func(int, int64, []byte) error {
				return nil
			})
			if err := checkUnusedNodesInvariant(tree); err != nil {
				t.Fatalf("UnusedNodesInvariant check failed: %v", err)
			}
		}
		if got, want := tree.Size(), int64(8); got != want {
			t.Errorf("Size()=%d, want %d", got, want)
		}
		if got, want := tree.CurrentRoot(), roots[7]; !bytes.Equal(got, want) {
			t.Errorf("CurrentRoot()=%v, want %v", got, want)
		}
		if diff := pretty.Compare(tree.Hashes(), hashes[7]); diff != "" {
			t.Errorf("post-Hashes() diff:\n%v", diff)
		}
	}
}

func failingGetNodeFunc(int, int64) ([]byte, error) {
	return []byte{}, errors.New("bang")
}

// This returns something that won't result in a valid root hash match, doesn't really
// matter what it is but it must be correct length for an SHA256 hash as if it was real
func fixedHashGetNodeFunc(int, int64) ([]byte, error) {
	return []byte("12345678901234567890123456789012"), nil
}

func TestLoadingTreeFailsNodeFetch(t *testing.T) {
	_, err := NewTreeWithState(rfc6962.DefaultHasher, 237, failingGetNodeFunc, []byte("notimportant"))

	if err == nil || !strings.Contains(err.Error(), "bang") {
		t.Errorf("Did not return correctly on failed node fetch: %v", err)
	}
}

func TestLoadingTreeFailsBadRootHash(t *testing.T) {
	// Supply a root hash that can't possibly match the result of the SHA 256 hashing on our dummy
	// data
	_, err := NewTreeWithState(rfc6962.DefaultHasher, 237, fixedHashGetNodeFunc, []byte("nomatch!nomatch!nomatch!nomatch!"))
	_, ok := err.(RootHashMismatchError)

	if err == nil || !ok {
		t.Errorf("Did not return correct error type on root mismatch: %v", err)
	}
}

func nodeKey(d int, i int64) (string, error) {
	n, err := storage.NewNodeIDForTreeCoords(int64(d), i, 64)
	if err != nil {
		return "", err
	}
	return n.String(), nil
}

func TestCompactVsFullTree(t *testing.T) {
	imt := merkle.NewInMemoryMerkleTree(rfc6962.DefaultHasher)
	nodes := make(map[string][]byte)

	for i := int64(0); i < 1024; i++ {
		cmt, err := NewTreeWithState(
			rfc6962.DefaultHasher,
			imt.LeafCount(),
			func(depth int, index int64) ([]byte, error) {
				k, err := nodeKey(depth, index)
				if err != nil {
					t.Errorf("failed to create nodeID: %v", err)
				}
				h := nodes[k]
				return h, nil
			}, imt.CurrentRoot().Hash())

		if err != nil {
			t.Errorf("interation %d: failed to create CMT with state: %v", i, err)
		}
		if a, b := imt.CurrentRoot().Hash(), cmt.CurrentRoot(); !bytes.Equal(a, b) {
			t.Errorf("iteration %d: Got in-memory root of %v, but compact tree has root %v", i, a, b)
		}

		newLeaf := []byte(fmt.Sprintf("Leaf %d", i))

		iSeq, iHash, err := imt.AddLeaf(newLeaf)
		if err != nil {
			t.Errorf("AddLeaf(): %v", err)
		}

		cSeq, cHash, err := cmt.AddLeaf(newLeaf,
			func(depth int, index int64, hash []byte) error {
				k, err := nodeKey(depth, index)
				if err != nil {
					return fmt.Errorf("failed to create nodeID: %v", err)
				}
				nodes[k] = hash
				return nil
			})
		if err != nil {
			t.Fatalf("mt update failed: %v", err)
		}

		// In-Memory tree is 1-based for sequence numbers, since it's based on the original CT C++ impl.
		if got, want := iSeq, i+1; got != want {
			t.Errorf("iteration %d: Got in-memory sequence number of %d, expected %d", i, got, want)
		}
		if int64(iSeq) != cSeq+1 {
			t.Errorf("iteration %d: Got in-memory sequence number of %d but %d (zero based) from compact tree", i, iSeq, cSeq)
		}
		if a, b := iHash.Hash(), cHash; !bytes.Equal(a, b) {
			t.Errorf("iteration %d: Got leaf hash %v from in-memory tree, but %v from compact tree", i, a, b)
		}
		if a, b := imt.CurrentRoot().Hash(), cmt.CurrentRoot(); !bytes.Equal(a, b) {
			t.Errorf("iteration %d: Got in-memory root of %v, but compact tree has root %v", i, a, b)
		}

	}
}

func TestRootHashForVariousTreeSizes(t *testing.T) {
	tests := []struct {
		size     int64
		wantRoot []byte
	}{
		{10, testonly.MustDecodeBase64("VjWMPSYNtCuCNlF/RLnQy6HcwSk6CIipfxm+hettA+4=")},
		{15, testonly.MustDecodeBase64("j4SulYmocFuxdeyp12xXCIgK6PekBcxzAIj4zbQzNEI=")},
		{16, testonly.MustDecodeBase64("c+4Uc6BCMOZf/v3NZK1kqTUJe+bBoFtOhP+P3SayKRE=")},
		{100, testonly.MustDecodeBase64("dUh9hYH88p0CMoHkdr1wC2szbhcLAXOejWpINIooKUY=")},
		{255, testonly.MustDecodeBase64("SmdsuKUqiod3RX2jyF2M6JnbdE4QuTwwipfAowI4/i0=")},
		{256, testonly.MustDecodeBase64("qFI0t/tZ1MdOYgyPpPzHFiZVw86koScXy9q3FU5casA=")},
		{1000, testonly.MustDecodeBase64("RXrgb8xHd55Y48FbfotJwCbV82Kx22LZfEbmBGAvwlQ=")},
		{4095, testonly.MustDecodeBase64("cWRFdQhPcjn9WyBXE/r1f04ejxIm5lvg40DEpRBVS0w=")},
		{4096, testonly.MustDecodeBase64("6uU/phfHg1n/GksYT6TO9aN8EauMCCJRl3dIK0HDs2M=")},
		{10000, testonly.MustDecodeBase64("VZcav65F9haHVRk3wre2axFoBXRNeUh/1d9d5FQfxIg=")},
		{65535, testonly.MustDecodeBase64("iPuVYJhP6SEE4gUFp8qbafd2rYv9YTCDYqAxCj8HdLM=")},
	}

	b64e := func(b []byte) string { return base64.StdEncoding.EncodeToString(b) }

	for _, test := range tests {
		tree := NewTree(rfc6962.DefaultHasher)
		for i := int64(0); i < test.size; i++ {
			l := []byte{byte(i & 0xff), byte((i >> 8) & 0xff)}
			tree.AddLeaf(l, func(int, int64, []byte) error {
				return nil
			})
		}
		if got, want := tree.CurrentRoot(), test.wantRoot; !bytes.Equal(got, want) {
			t.Errorf("Test (treesize=%v) got root %v, want %v", test.size, b64e(got), b64e(want))
		}
	}
}
<commit_msg>merkle/compact: Make TestAddingLeaves table-driven (#1342)<commit_after>// Copyright 2016 Google Inc. 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 compact

import (
	"bytes"
	"encoding/base64"
	"errors"
	"fmt"
	"math/bits"
	"strings"
	"testing"

	"github.com/google/trillian/merkle"
	"github.com/google/trillian/merkle/rfc6962"
	"github.com/google/trillian/storage"
	"github.com/google/trillian/testonly"
	"github.com/kylelemons/godebug/pretty"
)

// This check ensures that the compact Merkle tree contains the correct set of
// nodes, i.e. the node on level i is present iff i-th bit of tree size is 1.
func checkUnusedNodesInvariant(t *Tree) error {
	size := t.size
	sizeBits := bits.Len64(uint64(size))
	if got, want := len(t.nodes), sizeBits; got != want {
		return fmt.Errorf("nodes mismatch: have %v nodes, want %v", got, want)
	}
	for level := 0; level < sizeBits; level++ {
		if size&1 == 1 {
			if t.nodes[level] == nil {
				return fmt.Errorf("missing node at level %d", level)
			}
		} else if t.nodes[level] != nil {
			return fmt.Errorf("unexpected node at level %d", level)
		}
		size >>= 1
	}
	return nil
}

func TestAddingLeaves(t *testing.T) {
	inputs := testonly.MerkleTreeLeafTestInputs()
	roots := testonly.MerkleTreeLeafTestRootHashes()
	hashes := testonly.CompactMerkleTreeLeafTestNodeHashes()

	// Test the "same" thing in different ways, to ensure than any lazy update
	// strategy being employed by the implementation doesn't affect the
	// API-visible calculation of root & size.
	for _, tc := range []struct {
		desc   string
		breaks []int
	}{
		{desc: "one-by-one", breaks: []int{0, 1, 2, 3, 4, 5, 6, 7, 8}},
		{desc: "one-by-one-no-zero", breaks: []int{1, 2, 3, 4, 5, 6, 7, 8}},
		{desc: "all-at-once", breaks: []int{8}},
		{desc: "all-at-once-zero", breaks: []int{0, 8}},
		{desc: "two-chunks", breaks: []int{3, 8}},
		{desc: "two-chunks-zero", breaks: []int{0, 3, 8}},
	} {
		t.Run(tc.desc, func(t *testing.T) {
			tree := NewTree(rfc6962.DefaultHasher)
			idx := 0
			for _, br := range tc.breaks {
				for ; idx < br; idx++ {
					if _, _, err := tree.AddLeaf(inputs[idx], func(int, int64, []byte) error {
						return nil
					}); err != nil {
						t.Fatalf("AddLeaf: %v", err)
					}
					if err := checkUnusedNodesInvariant(tree); err != nil {
						t.Fatalf("UnusedNodesInvariant check failed: %v", err)
					}
				}
				if got, want := tree.Size(), int64(br); got != want {
					t.Errorf("Size()=%d, want %d", got, want)
				}
				if br > 0 {
					if got, want := tree.CurrentRoot(), roots[br-1]; !bytes.Equal(got, want) {
						t.Errorf("CurrentRoot()=%v, want %v", got, want)
					}
					if diff := pretty.Compare(tree.Hashes(), hashes[br-1]); diff != "" {
						t.Errorf("post-Hashes() diff:\n%v", diff)
					}
				} else {
					if got, want := tree.CurrentRoot(), testonly.EmptyMerkleTreeRootHash(); !bytes.Equal(got, want) {
						t.Errorf("CurrentRoot()=%x, want %x (empty)", got, want)
					}
				}
			}
		})
	}
}

func failingGetNodeFunc(int, int64) ([]byte, error) {
	return []byte{}, errors.New("bang")
}

// This returns something that won't result in a valid root hash match, doesn't really
// matter what it is but it must be correct length for an SHA256 hash as if it was real
func fixedHashGetNodeFunc(int, int64) ([]byte, error) {
	return []byte("12345678901234567890123456789012"), nil
}

func TestLoadingTreeFailsNodeFetch(t *testing.T) {
	_, err := NewTreeWithState(rfc6962.DefaultHasher, 237, failingGetNodeFunc, []byte("notimportant"))

	if err == nil || !strings.Contains(err.Error(), "bang") {
		t.Errorf("Did not return correctly on failed node fetch: %v", err)
	}
}

func TestLoadingTreeFailsBadRootHash(t *testing.T) {
	// Supply a root hash that can't possibly match the result of the SHA 256 hashing on our dummy
	// data
	_, err := NewTreeWithState(rfc6962.DefaultHasher, 237, fixedHashGetNodeFunc, []byte("nomatch!nomatch!nomatch!nomatch!"))
	_, ok := err.(RootHashMismatchError)

	if err == nil || !ok {
		t.Errorf("Did not return correct error type on root mismatch: %v", err)
	}
}

func nodeKey(d int, i int64) (string, error) {
	n, err := storage.NewNodeIDForTreeCoords(int64(d), i, 64)
	if err != nil {
		return "", err
	}
	return n.String(), nil
}

func TestCompactVsFullTree(t *testing.T) {
	imt := merkle.NewInMemoryMerkleTree(rfc6962.DefaultHasher)
	nodes := make(map[string][]byte)

	for i := int64(0); i < 1024; i++ {
		cmt, err := NewTreeWithState(
			rfc6962.DefaultHasher,
			imt.LeafCount(),
			func(depth int, index int64) ([]byte, error) {
				k, err := nodeKey(depth, index)
				if err != nil {
					t.Errorf("failed to create nodeID: %v", err)
				}
				h := nodes[k]
				return h, nil
			}, imt.CurrentRoot().Hash())

		if err != nil {
			t.Errorf("interation %d: failed to create CMT with state: %v", i, err)
		}
		if a, b := imt.CurrentRoot().Hash(), cmt.CurrentRoot(); !bytes.Equal(a, b) {
			t.Errorf("iteration %d: Got in-memory root of %v, but compact tree has root %v", i, a, b)
		}

		newLeaf := []byte(fmt.Sprintf("Leaf %d", i))

		iSeq, iHash, err := imt.AddLeaf(newLeaf)
		if err != nil {
			t.Errorf("AddLeaf(): %v", err)
		}

		cSeq, cHash, err := cmt.AddLeaf(newLeaf,
			func(depth int, index int64, hash []byte) error {
				k, err := nodeKey(depth, index)
				if err != nil {
					return fmt.Errorf("failed to create nodeID: %v", err)
				}
				nodes[k] = hash
				return nil
			})
		if err != nil {
			t.Fatalf("mt update failed: %v", err)
		}

		// In-Memory tree is 1-based for sequence numbers, since it's based on the original CT C++ impl.
		if got, want := iSeq, i+1; got != want {
			t.Errorf("iteration %d: Got in-memory sequence number of %d, expected %d", i, got, want)
		}
		if int64(iSeq) != cSeq+1 {
			t.Errorf("iteration %d: Got in-memory sequence number of %d but %d (zero based) from compact tree", i, iSeq, cSeq)
		}
		if a, b := iHash.Hash(), cHash; !bytes.Equal(a, b) {
			t.Errorf("iteration %d: Got leaf hash %v from in-memory tree, but %v from compact tree", i, a, b)
		}
		if a, b := imt.CurrentRoot().Hash(), cmt.CurrentRoot(); !bytes.Equal(a, b) {
			t.Errorf("iteration %d: Got in-memory root of %v, but compact tree has root %v", i, a, b)
		}

	}
}

func TestRootHashForVariousTreeSizes(t *testing.T) {
	tests := []struct {
		size     int64
		wantRoot []byte
	}{
		{10, testonly.MustDecodeBase64("VjWMPSYNtCuCNlF/RLnQy6HcwSk6CIipfxm+hettA+4=")},
		{15, testonly.MustDecodeBase64("j4SulYmocFuxdeyp12xXCIgK6PekBcxzAIj4zbQzNEI=")},
		{16, testonly.MustDecodeBase64("c+4Uc6BCMOZf/v3NZK1kqTUJe+bBoFtOhP+P3SayKRE=")},
		{100, testonly.MustDecodeBase64("dUh9hYH88p0CMoHkdr1wC2szbhcLAXOejWpINIooKUY=")},
		{255, testonly.MustDecodeBase64("SmdsuKUqiod3RX2jyF2M6JnbdE4QuTwwipfAowI4/i0=")},
		{256, testonly.MustDecodeBase64("qFI0t/tZ1MdOYgyPpPzHFiZVw86koScXy9q3FU5casA=")},
		{1000, testonly.MustDecodeBase64("RXrgb8xHd55Y48FbfotJwCbV82Kx22LZfEbmBGAvwlQ=")},
		{4095, testonly.MustDecodeBase64("cWRFdQhPcjn9WyBXE/r1f04ejxIm5lvg40DEpRBVS0w=")},
		{4096, testonly.MustDecodeBase64("6uU/phfHg1n/GksYT6TO9aN8EauMCCJRl3dIK0HDs2M=")},
		{10000, testonly.MustDecodeBase64("VZcav65F9haHVRk3wre2axFoBXRNeUh/1d9d5FQfxIg=")},
		{65535, testonly.MustDecodeBase64("iPuVYJhP6SEE4gUFp8qbafd2rYv9YTCDYqAxCj8HdLM=")},
	}

	b64e := func(b []byte) string { return base64.StdEncoding.EncodeToString(b) }

	for _, test := range tests {
		tree := NewTree(rfc6962.DefaultHasher)
		for i := int64(0); i < test.size; i++ {
			l := []byte{byte(i & 0xff), byte((i >> 8) & 0xff)}
			tree.AddLeaf(l, func(int, int64, []byte) error {
				return nil
			})
		}
		if got, want := tree.CurrentRoot(), test.wantRoot; !bytes.Equal(got, want) {
			t.Errorf("Test (treesize=%v) got root %v, want %v", test.size, b64e(got), b64e(want))
		}
	}
}
<|endoftext|>


<commit_before>// Package aws provides core functionality for making requests to AWS services.
package aws

// SDKName is the name of this AWS SDK
const SDKName = "aws-sdk-go"

// SDKVersion is the version of this SDK
const SDKVersion = "1.1.9"
<commit_msg>Tag release v1.1.10<commit_after>// Package aws provides core functionality for making requests to AWS services.
package aws

// SDKName is the name of this AWS SDK
const SDKName = "aws-sdk-go"

// SDKVersion is the version of this SDK
const SDKVersion = "1.1.10"
<|endoftext|>


<commit_before>/*
 * Copyright 2015 Google Inc. 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.
 */

// Binary kwazthis (K, what's this?) determines what references are located at a
// particular offset within a file.  All results are printed as JSON.
//
// By default, kwazthis will search for a .kythe configuration file in a
// directory above the given --path (if it exists locally relative to the
// current working directory).  If found, --path will be made relative to this
// directory and --root before making any Kythe service requests.  If not found,
// --path will be passed unchanged.  --ignore_local_repo will turn off this
// behavior.
//
// Usage:
//   kwazthis --path kythe/cxx/tools/kindex_tool_main.cc --offset 2660
//   kwazthis --path kythe/java/com/google/devtools/kythe/analyzers/base/EntrySet.java --offset 2815
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"kythe.io/kythe/go/services/search"
	"kythe.io/kythe/go/services/xrefs"
	"kythe.io/kythe/go/util/kytheuri"
	"kythe.io/kythe/go/util/schema"

	"golang.org/x/net/context"
	"google.golang.org/grpc"

	spb "kythe.io/kythe/proto/storage_proto"
	xpb "kythe.io/kythe/proto/xref_proto"
)

func init() {
	binary := filepath.Base(os.Args[0])
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, `Determine what references are located at a particular offset within a file.

Usage: %s --offset int (--path p | --signature s) [--corpus c] [--root r] [--language l]
       %`+strconv.Itoa(len(binary))+`s [--ignore_local_repo] [--dirty_buffer path]

By default, kwazthis will search for a .kythe configuration file in a directory
above the given --path (if it exists locally relative to the current working
directory).  If found, --path will be made relative to this directory and --root
before making any Kythe service requests.  If not found, --path will be passed
unchanged.  --ignore_local_repo will turn off this behavior.

Defaults flag values:
`, binary, "")
		flag.PrintDefaults()
	}
}

var (
	remoteAPI = flag.String("api", "https://xrefs-dot-kythe-repo.appspot.com", "Remote API server")

	ignoreLocalRepo = flag.Bool("ignore_local_repo", false, "Ignore local repository .kythe configuration")

	dirtyBuffer = flag.String("dirty_buffer", "", "Path to file with dirty buffer contents (optional)")

	path      = flag.String("path", "", "Path of file (optional if --signature is given)")
	signature = flag.String("signature", "", "Signature of file VName (optional if --path is given)")
	corpus    = flag.String("corpus", "", "Corpus of file VName (optional)")
	root      = flag.String("root", "", "Root of file VName (optional)")
	language  = flag.String("language", "", "Language of file VName (optional)")

	offset = flag.Int("offset", -1, "Non-negative offset in file to list references")
)

var (
	xs  xrefs.Service
	idx search.Service

	fileFacts = []*spb.SearchRequest_Fact{
		{Name: schema.NodeKindFact, Value: []byte(schema.FileKind)},
	}
)

type reference struct {
	Span struct {
		Start int    `json:"start"`
		End   int    `json:"end"`
		Text  string `json:"text,omitempty"`
	} `json:"span"`
	Kind string `json:"kind"`

	Node struct {
		Ticket  string   `json:"ticket"`
		Names   []string `json:"names,omitempty"`
		Kind    string   `json:"kind,omitempty"`
		Subkind string   `json:"subkind,omitempty"`
	} `json:"node"`
}

func main() {
	flag.Parse()
	if *offset < 0 {
		log.Fatal("ERROR: non-negative --offset required")
	} else if *signature == "" && *path == "" {
		log.Fatal("ERROR: must provide at least -path or --signature")
	}

	if strings.HasPrefix(*remoteAPI, "http://") || strings.HasPrefix(*remoteAPI, "https://") {
		xs = xrefs.WebClient(*remoteAPI)
		idx = search.WebClient(*remoteAPI)
	} else {
		conn, err := grpc.Dial(*remoteAPI)
		if err != nil {
			log.Fatalf("Error connecting to remote API %q: %v", *remoteAPI, err)
		}
		defer conn.Close()
		ctx := context.Background()
		xs = xrefs.GRPC(ctx, xpb.NewXRefServiceClient(conn))
		idx = search.GRPC(ctx, spb.NewSearchServiceClient(conn))
	}

	relPath := *path
	if !*ignoreLocalRepo {
		if _, err := os.Stat(relPath); err == nil {
			absPath, err := filepath.Abs(relPath)
			if err != nil {
				log.Fatal(err)
			}
			kytheRoot := findKytheRoot(filepath.Dir(absPath))
			if kytheRoot != "" {
				relPath, err = filepath.Rel(filepath.Join(kytheRoot, *root), absPath)
				if err != nil {
					log.Fatal(err)
				}
			}
		}
	}
	partialFile := &spb.VName{
		Signature: *signature,
		Corpus:    *corpus,
		Root:      *root,
		Path:      relPath,
		Language:  *language,
	}
	reply, err := idx.Search(&spb.SearchRequest{
		Partial: partialFile,
		Fact:    fileFacts,
	})
	if err != nil {
		log.Fatalf("Error locating file {%v}: %v", partialFile, err)
	}
	if len(reply.Ticket) == 0 {
		log.Fatalf("Could not locate file {%v}", partialFile)
	} else if len(reply.Ticket) > 1 {
		log.Fatalf("Ambiguous file {%v}; multiple results: %v", partialFile, reply.Ticket)
	}

	fileTicket := reply.Ticket[0]
	decor, err := xs.Decorations(&xpb.DecorationsRequest{
		// TODO(schroederc): limit Location to a SPAN around *offset
		Location:    &xpb.Location{Ticket: fileTicket},
		References:  true,
		SourceText:  true,
		DirtyBuffer: readDirtyBuffer(),
	})
	if err != nil {
		log.Fatal(err)
	}
	nodes := xrefs.NodesMap(decor.Node)

	en := json.NewEncoder(os.Stdout)
	for _, ref := range decor.Reference {
		start, _ := strconv.Atoi(string(nodes[ref.SourceTicket][schema.AnchorStartFact]))
		end, _ := strconv.Atoi(string(nodes[ref.SourceTicket][schema.AnchorEndFact]))

		if start <= *offset && *offset < end {
			var r reference
			r.Span.Start = start
			r.Span.End = end
			r.Span.Text = string(decor.SourceText[start:end])
			r.Kind = strings.TrimPrefix(ref.Kind, schema.EdgePrefix)
			r.Node.Ticket = ref.TargetTicket

			node := nodes[ref.TargetTicket]
			r.Node.Kind = string(node[schema.NodeKindFact])
			r.Node.Subkind = string(node[schema.SubkindFact])

			if eReply, err := xs.Edges(&xpb.EdgesRequest{
				Ticket: []string{ref.TargetTicket},
				Kind:   []string{schema.NamedEdge},
			}); err != nil {
				log.Println("WARNING: error getting edges for %q: %v", ref.TargetTicket, err)
			} else {
				for _, name := range xrefs.EdgesMap(eReply.EdgeSet)[ref.TargetTicket][schema.NamedEdge] {
					if uri, err := kytheuri.Parse(name); err != nil {
						log.Println("WARNING: named node ticket (%q) could not be parsed: %v", name, err)
					} else {
						r.Node.Names = append(r.Node.Names, uri.Signature)
					}
				}
			}

			if err := en.Encode(r); err != nil {
				log.Fatal(err)
			}
		}
	}
}

func readDirtyBuffer() []byte {
	if *dirtyBuffer == "" {
		return nil
	}

	f, err := os.Open(*dirtyBuffer)
	if err != nil {
		log.Fatal("ERROR: could not open dirty buffer at %q: %v", *dirtyBuffer, err)
	}
	defer f.Close()
	data, err := ioutil.ReadAll(f)
	if err != nil {
		log.Fatal("ERROR: could read dirty buffer at %q: %v", *dirtyBuffer, err)
	}
	return data
}

func findKytheRoot(dir string) string {
	for {
		if fi, err := os.Stat(filepath.Join(dir, ".kythe")); err == nil && fi.Mode().IsRegular() {
			return dir
		}
		if dir == "/" {
			break
		}
		dir = filepath.Dir(dir)
	}
	return ""
}
<commit_msg>Fix some Go logging statements<commit_after>/*
 * Copyright 2015 Google Inc. 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.
 */

// Binary kwazthis (K, what's this?) determines what references are located at a
// particular offset within a file.  All results are printed as JSON.
//
// By default, kwazthis will search for a .kythe configuration file in a
// directory above the given --path (if it exists locally relative to the
// current working directory).  If found, --path will be made relative to this
// directory and --root before making any Kythe service requests.  If not found,
// --path will be passed unchanged.  --ignore_local_repo will turn off this
// behavior.
//
// Usage:
//   kwazthis --path kythe/cxx/tools/kindex_tool_main.cc --offset 2660
//   kwazthis --path kythe/java/com/google/devtools/kythe/analyzers/base/EntrySet.java --offset 2815
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"kythe.io/kythe/go/services/search"
	"kythe.io/kythe/go/services/xrefs"
	"kythe.io/kythe/go/util/kytheuri"
	"kythe.io/kythe/go/util/schema"

	"golang.org/x/net/context"
	"google.golang.org/grpc"

	spb "kythe.io/kythe/proto/storage_proto"
	xpb "kythe.io/kythe/proto/xref_proto"
)

func init() {
	binary := filepath.Base(os.Args[0])
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, `Determine what references are located at a particular offset within a file.

Usage: %s --offset int (--path p | --signature s) [--corpus c] [--root r] [--language l]
       %`+strconv.Itoa(len(binary))+`s [--ignore_local_repo] [--dirty_buffer path]

By default, kwazthis will search for a .kythe configuration file in a directory
above the given --path (if it exists locally relative to the current working
directory).  If found, --path will be made relative to this directory and --root
before making any Kythe service requests.  If not found, --path will be passed
unchanged.  --ignore_local_repo will turn off this behavior.

Defaults flag values:
`, binary, "")
		flag.PrintDefaults()
	}
}

var (
	remoteAPI = flag.String("api", "https://xrefs-dot-kythe-repo.appspot.com", "Remote API server")

	ignoreLocalRepo = flag.Bool("ignore_local_repo", false, "Ignore local repository .kythe configuration")

	dirtyBuffer = flag.String("dirty_buffer", "", "Path to file with dirty buffer contents (optional)")

	path      = flag.String("path", "", "Path of file (optional if --signature is given)")
	signature = flag.String("signature", "", "Signature of file VName (optional if --path is given)")
	corpus    = flag.String("corpus", "", "Corpus of file VName (optional)")
	root      = flag.String("root", "", "Root of file VName (optional)")
	language  = flag.String("language", "", "Language of file VName (optional)")

	offset = flag.Int("offset", -1, "Non-negative offset in file to list references")
)

var (
	xs  xrefs.Service
	idx search.Service

	fileFacts = []*spb.SearchRequest_Fact{
		{Name: schema.NodeKindFact, Value: []byte(schema.FileKind)},
	}
)

type reference struct {
	Span struct {
		Start int    `json:"start"`
		End   int    `json:"end"`
		Text  string `json:"text,omitempty"`
	} `json:"span"`
	Kind string `json:"kind"`

	Node struct {
		Ticket  string   `json:"ticket"`
		Names   []string `json:"names,omitempty"`
		Kind    string   `json:"kind,omitempty"`
		Subkind string   `json:"subkind,omitempty"`
	} `json:"node"`
}

func main() {
	flag.Parse()
	if *offset < 0 {
		log.Fatal("ERROR: non-negative --offset required")
	} else if *signature == "" && *path == "" {
		log.Fatal("ERROR: must provide at least -path or --signature")
	}

	if strings.HasPrefix(*remoteAPI, "http://") || strings.HasPrefix(*remoteAPI, "https://") {
		xs = xrefs.WebClient(*remoteAPI)
		idx = search.WebClient(*remoteAPI)
	} else {
		conn, err := grpc.Dial(*remoteAPI)
		if err != nil {
			log.Fatalf("Error connecting to remote API %q: %v", *remoteAPI, err)
		}
		defer conn.Close()
		ctx := context.Background()
		xs = xrefs.GRPC(ctx, xpb.NewXRefServiceClient(conn))
		idx = search.GRPC(ctx, spb.NewSearchServiceClient(conn))
	}

	relPath := *path
	if !*ignoreLocalRepo {
		if _, err := os.Stat(relPath); err == nil {
			absPath, err := filepath.Abs(relPath)
			if err != nil {
				log.Fatal(err)
			}
			kytheRoot := findKytheRoot(filepath.Dir(absPath))
			if kytheRoot != "" {
				relPath, err = filepath.Rel(filepath.Join(kytheRoot, *root), absPath)
				if err != nil {
					log.Fatal(err)
				}
			}
		}
	}
	partialFile := &spb.VName{
		Signature: *signature,
		Corpus:    *corpus,
		Root:      *root,
		Path:      relPath,
		Language:  *language,
	}
	reply, err := idx.Search(&spb.SearchRequest{
		Partial: partialFile,
		Fact:    fileFacts,
	})
	if err != nil {
		log.Fatalf("Error locating file {%v}: %v", partialFile, err)
	}
	if len(reply.Ticket) == 0 {
		log.Fatalf("Could not locate file {%v}", partialFile)
	} else if len(reply.Ticket) > 1 {
		log.Fatalf("Ambiguous file {%v}; multiple results: %v", partialFile, reply.Ticket)
	}

	fileTicket := reply.Ticket[0]
	decor, err := xs.Decorations(&xpb.DecorationsRequest{
		// TODO(schroederc): limit Location to a SPAN around *offset
		Location:    &xpb.Location{Ticket: fileTicket},
		References:  true,
		SourceText:  true,
		DirtyBuffer: readDirtyBuffer(),
	})
	if err != nil {
		log.Fatal(err)
	}
	nodes := xrefs.NodesMap(decor.Node)

	en := json.NewEncoder(os.Stdout)
	for _, ref := range decor.Reference {
		start, _ := strconv.Atoi(string(nodes[ref.SourceTicket][schema.AnchorStartFact]))
		end, _ := strconv.Atoi(string(nodes[ref.SourceTicket][schema.AnchorEndFact]))

		if start <= *offset && *offset < end {
			var r reference
			r.Span.Start = start
			r.Span.End = end
			r.Span.Text = string(decor.SourceText[start:end])
			r.Kind = strings.TrimPrefix(ref.Kind, schema.EdgePrefix)
			r.Node.Ticket = ref.TargetTicket

			node := nodes[ref.TargetTicket]
			r.Node.Kind = string(node[schema.NodeKindFact])
			r.Node.Subkind = string(node[schema.SubkindFact])

			if eReply, err := xs.Edges(&xpb.EdgesRequest{
				Ticket: []string{ref.TargetTicket},
				Kind:   []string{schema.NamedEdge},
			}); err != nil {
				log.Printf("WARNING: error getting edges for %q: %v", ref.TargetTicket, err)
			} else {
				for _, name := range xrefs.EdgesMap(eReply.EdgeSet)[ref.TargetTicket][schema.NamedEdge] {
					if uri, err := kytheuri.Parse(name); err != nil {
						log.Printf("WARNING: named node ticket (%q) could not be parsed: %v", name, err)
					} else {
						r.Node.Names = append(r.Node.Names, uri.Signature)
					}
				}
			}

			if err := en.Encode(r); err != nil {
				log.Fatal(err)
			}
		}
	}
}

func readDirtyBuffer() []byte {
	if *dirtyBuffer == "" {
		return nil
	}

	f, err := os.Open(*dirtyBuffer)
	if err != nil {
		log.Fatalf("ERROR: could not open dirty buffer at %q: %v", *dirtyBuffer, err)
	}
	defer f.Close()
	data, err := ioutil.ReadAll(f)
	if err != nil {
		log.Fatalf("ERROR: could read dirty buffer at %q: %v", *dirtyBuffer, err)
	}
	return data
}

func findKytheRoot(dir string) string {
	for {
		if fi, err := os.Stat(filepath.Join(dir, ".kythe")); err == nil && fi.Mode().IsRegular() {
			return dir
		}
		if dir == "/" {
			break
		}
		dir = filepath.Dir(dir)
	}
	return ""
}
<|endoftext|>


<commit_before>package base62

import (
	"fmt"
	"math/big"
	"sort"
	"strconv"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

var result string

var testcases = []struct {
	num     int64
	encoded string
}{
	{1, "1"},
	{9, "9"},
	{10, "A"},
	{35, "Z"},
	{36, "a"},
	{61, "z"},
	{62, "10"},
	{99, "1b"},
	{3844, "100"},
	{3860, "10G"},
	{4815162342, "5Frvgk"},
	{9223372036854775807, "AzL8n0Y58m7"},
}

func TestEncodeInt64(t *testing.T) {
	for _, tc := range testcases {
		v := EncodeInt64(tc.num)
		t.Logf("Encoded %v as %s", tc.num, v)
		assert.Equal(t, tc.encoded, v)
	}
}

func TestDecodeToInt64(t *testing.T) {
	for _, tc := range testcases {
		v := DecodeToInt64(tc.encoded)
		t.Logf("Decoded %s to %v", tc.encoded, v)
		assert.Equal(t, tc.num, v)
	}
}

func BenchmarkEncodeInt64Medium(b *testing.B) {
	var id string
	for n := 0; n < b.N; n++ {
		id = EncodeInt64(4815162342)
	}
	result = id
}

func BenchmarkEncodeInt64Long(b *testing.B) {
	var id string
	for n := 0; n < b.N; n++ {
		id = EncodeInt64(9223372036854775807)
	}
	result = id
}

var bigTestcases = []struct {
	num     string
	encoded string
}{
	{"1", "1"},
	{"9", "9"},
	{"10", "A"},
	{"35", "Z"},
	{"36", "a"},
	{"61", "z"},
	{"62", "10"},
	{"99", "1b"},
	{"3844", "100"},
	{"3860", "10G"},
	{"4815162342", "5Frvgk"},

	{"9223372036854775807", "AzL8n0Y58m7"},       // max signed int64
	{"9223372036854775809", "AzL8n0Y58m9"},       // beyond int64
	{"9223372036854775861", "AzL8n0Y58mz"},       //
	{"18446744073709551615", "LygHa16AHYF"},      // max uint64
	{"571849066284996100034", "AzL8n0Y58m70"},    // max int64 * 62
	{"35454642109669758202168", "AzL8n0Y58m70y"}, // (max int64 * 62^2) + 60

	{"170141183460469231731687303715884105727", "3tX16dB2jpss4tZORYcqo3"}, // max signed 128bit int
	{"170141183460469231731687303715884105757", "3tX16dB2jpss4tZORYcqoX"}, // max signed 128bit int + 30
	{"340282366920938463463374607431768211455", "7n42DGM5Tflk9n8mt7Fhc7"}, // max unsigned 128bit int
}

func TestEncodeBigInt(t *testing.T) {
	for _, tc := range bigTestcases {
		var (
			n  = new(big.Int)
			ok bool
		)

		n, ok = n.SetString(tc.num, 10)
		require.True(t, ok)

		v := EncodeBigInt(n)
		t.Logf("Encoded %v as %s", tc.num, v)
		assert.Equal(t, tc.encoded, v)
	}
}

func BenchmarkEncodeBigIntVeryLong(b *testing.B) {
	var (
		v = new(big.Int)
		s string
	)
	v.SetString("340282366920938463463374607431768211455", 10)

	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		s = EncodeBigInt(v)
	}
	result = s
}

func TestDecodeToBigInt(t *testing.T) {
	for _, tc := range bigTestcases {
		v := DecodeToBigInt(tc.encoded)
		t.Logf("Decoded %v to %s", tc.encoded, v.String())
		assert.Equal(t, tc.num, v.String())
	}
}

// TestLexicalPaddedSort tests that numbers encoded as base62 strings
// are correctly lexically sorted with the original order preserved
// if these are left padded to the same length.
//
// An alternative sort method which could be used to avoid padding
// would be a Shortlex sort, which sorts by cardinality, then lexically.
func TestLexicalPaddedSort(t *testing.T) {

	var (
		lexicalOrder  sort.StringSlice = make([]string, 0)
		originalOrder                  = make([]string, 0)
	)

	// Generate lots of numbers, and encode them
	var i int64
	var modifier int64 = 1
	for i = 0; i < 100000; i++ {
		if i%10000 == 0 {
			modifier = modifier * 30
		}

		v := EncodeInt64(i + modifier)

		lexicalOrder = append(lexicalOrder, v)
		originalOrder = append(originalOrder, v)
	}

	// Find longest string & pad encoded strings to this length
	maxlen := len(originalOrder[len(originalOrder)-1])
	originalOrder = padStringArray(originalOrder, maxlen)
	lexicalOrder = padStringArray(lexicalOrder, maxlen)

	// Sort string array
	lexicalOrder.Sort()

	// Compare ordering with original
	var mismatch int64
	for i, v := range originalOrder {
		// t.Logf("%s %s", v, lexicalOrder[i])
		if lexicalOrder[i] != v {
			mismatch++
		}
	}
	assert.Equal(t, int64(0), mismatch, fmt.Sprintf("Expected zero mismatches, got %v", mismatch))
}

func padStringArray(s []string, maxlen int) []string {

	for i, v := range s {
		s[i] = pad(v, maxlen)
	}

	return s
}

func pad(s string, maxlen int) string {
	format := fmt.Sprint(`%0`, strconv.Itoa(maxlen), "s")
	return fmt.Sprintf(format, s)
}
<commit_msg>Add a few more test cases<commit_after>package base62

import (
	"fmt"
	"math/big"
	"sort"
	"strconv"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

var result string

var testcases = []struct {
	num     int64
	encoded string
}{
	{1, "1"},
	{9, "9"},
	{10, "A"},
	{35, "Z"},
	{36, "a"},
	{61, "z"},
	{62, "10"},
	{99, "1b"},
	{3844, "100"},
	{3860, "10G"},
	{4815162342, "5Frvgk"},
	{9223372036854775807, "AzL8n0Y58m7"},
}

func TestEncodeInt64(t *testing.T) {
	for _, tc := range testcases {
		v := EncodeInt64(tc.num)
		t.Logf("Encoded %v as %s", tc.num, v)
		assert.Equal(t, tc.encoded, v)
	}
}

func TestDecodeToInt64(t *testing.T) {
	for _, tc := range testcases {
		v := DecodeToInt64(tc.encoded)
		t.Logf("Decoded %s to %v", tc.encoded, v)
		assert.Equal(t, tc.num, v)
	}
}

func BenchmarkEncodeInt64Medium(b *testing.B) {
	var id string
	for n := 0; n < b.N; n++ {
		id = EncodeInt64(4815162342)
	}
	result = id
}

func BenchmarkEncodeInt64Long(b *testing.B) {
	var id string
	for n := 0; n < b.N; n++ {
		id = EncodeInt64(9223372036854775807)
	}
	result = id
}

var bigTestcases = []struct {
	num     string
	encoded string
}{
	{"1", "1"},
	{"9", "9"},
	{"10", "A"},
	{"35", "Z"},
	{"36", "a"},
	{"61", "z"},
	{"62", "10"},
	{"99", "1b"},
	{"3844", "100"},
	{"3860", "10G"},
	{"4815162342", "5Frvgk"},

	{"9223372036854775807", "AzL8n0Y58m7"},       // max signed int64
	{"9223372036854775809", "AzL8n0Y58m9"},       // beyond int64
	{"9223372036854775861", "AzL8n0Y58mz"},       //
	{"18446744073709551615", "LygHa16AHYF"},      // max uint64
	{"571849066284996100034", "AzL8n0Y58m70"},    // max int64 * 62
	{"35454642109669758202168", "AzL8n0Y58m70y"}, // (max int64 * 62^2) + 60

	{"24467927614188555520896788267013", "8HFaR8qWtRlGDHnO57"}, // a few boundary flake id tests
	{"24467927614170108776823078715395", "8HFaR8qAulTgCBd6Wp"},
	{"24467927614170108776823078715394", "8HFaR8qAulTgCBd6Wo"},
	{"24467927614170108776823078715393", "8HFaR8qAulTgCBd6Wn"},
	{"24467927614170108776823078715392", "8HFaR8qAulTgCBd6Wm"},

	{"170141183460469231731687303715884105727", "3tX16dB2jpss4tZORYcqo3"}, // max signed 128bit int
	{"170141183460469231731687303715884105757", "3tX16dB2jpss4tZORYcqoX"}, // max signed 128bit int + 30
	{"340282366920938463463374607431768211455", "7n42DGM5Tflk9n8mt7Fhc7"}, // max unsigned 128bit int
}

func TestEncodeBigInt(t *testing.T) {
	for _, tc := range bigTestcases {
		var (
			n  = new(big.Int)
			ok bool
		)

		n, ok = n.SetString(tc.num, 10)
		require.True(t, ok)

		v := EncodeBigInt(n)
		t.Logf("Encoded %v as %s", tc.num, v)
		assert.Equal(t, tc.encoded, v)
	}
}

func BenchmarkEncodeBigIntVeryLong(b *testing.B) {
	var (
		v = new(big.Int)
		s string
	)
	v.SetString("340282366920938463463374607431768211455", 10)

	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		s = EncodeBigInt(v)
	}
	result = s
}

func TestDecodeToBigInt(t *testing.T) {
	for _, tc := range bigTestcases {
		v := DecodeToBigInt(tc.encoded)
		t.Logf("Decoded %v to %s", tc.encoded, v.String())
		assert.Equal(t, tc.num, v.String())
	}
}

// TestLexicalPaddedSort tests that numbers encoded as base62 strings
// are correctly lexically sorted with the original order preserved
// if these are left padded to the same length.
//
// An alternative sort method which could be used to avoid padding
// would be a Shortlex sort, which sorts by cardinality, then lexically.
func TestLexicalPaddedSort(t *testing.T) {

	var (
		lexicalOrder  sort.StringSlice = make([]string, 0)
		originalOrder                  = make([]string, 0)
	)

	// Generate lots of numbers, and encode them
	var i int64
	var modifier int64 = 1
	for i = 0; i < 100000; i++ {
		if i%10000 == 0 {
			modifier = modifier * 30
		}

		v := EncodeInt64(i + modifier)

		lexicalOrder = append(lexicalOrder, v)
		originalOrder = append(originalOrder, v)
	}

	// Find longest string & pad encoded strings to this length
	maxlen := len(originalOrder[len(originalOrder)-1])
	originalOrder = padStringArray(originalOrder, maxlen)
	lexicalOrder = padStringArray(lexicalOrder, maxlen)

	// Sort string array
	lexicalOrder.Sort()

	// Compare ordering with original
	var mismatch int64
	for i, v := range originalOrder {
		// t.Logf("%s %s", v, lexicalOrder[i])
		if lexicalOrder[i] != v {
			mismatch++
		}
	}
	assert.Equal(t, int64(0), mismatch, fmt.Sprintf("Expected zero mismatches, got %v", mismatch))
}

func padStringArray(s []string, maxlen int) []string {

	for i, v := range s {
		s[i] = pad(v, maxlen)
	}

	return s
}

func pad(s string, maxlen int) string {
	format := fmt.Sprint(`%0`, strconv.Itoa(maxlen), "s")
	return fmt.Sprintf(format, s)
}
<|endoftext|>