<commit_before>package catalog

import (
	"fmt"
	"log"
	"sync"

	"github.com/dnaeon/gru/graph"
	"github.com/dnaeon/gru/resource"
	"github.com/dnaeon/gru/utils"
	"github.com/layeh/gopher-luar"
	"github.com/yuin/gopher-lua"
)

// Catalog type contains a collection of resources
type Catalog struct {
	// Unsorted contains the list of resources created by Lua
	Unsorted []resource.Resource `luar:"-"`

	// Collection contains the unsorted resources as a collection
	collection resource.Collection `luar:"-"`

	// Sorted contains the resources after a topological sort.
	sorted []*graph.Node `luar:"-"`

	// Reversed contains the resource dependency graph in reverse
	// order. It is used for finding the reverse dependencies of
	// resources.
	reversed *graph.Graph `luar:"-"`

	// Status contains status information about resources
	status *status `luar:"-"`

	// Configuration settings
	config *Config `luar:"-"`
}

// Config type represents a set of settings to use when
// creating and processing the catalog
type Config struct {
	// Name of the Lua module to load and execute
	Module string

	// Do not take any actions, just report what would be done
	DryRun bool

	// Writer used to log events
	Logger *log.Logger

	// Path to the site repo containing module and data files
	SiteRepo string

	// The Lua state
	L *lua.LState

	// Number of goroutines to use for concurrent processing
	Concurrency int
}

// status type contains status information about processed resources
type status struct {
	sync.RWMutex

	// Items contain the result of resource processing and any
	// errors that might have occurred during processing.
	// Keys of the map are the resource ids and their
	// values are the errors returned by resources.
	items map[string]error
}

// set sets the status for a resource
func (s *status) set(id string, err error) {
	s.Lock()
	defer s.Unlock()
	s.items[id] = err
}

// get retrieves the status of a resource
func (s *status) get(id string) (error, bool) {
	s.Lock()
	defer s.Unlock()
	err, ok := s.items[id]

	return err, ok
}

// isSynced returns a boolean indicating whether a
// resource is up to date
func (s *status) isSynced(id string) bool {
	s.Lock()
	defer s.Unlock()

	return s.items[id] == resource.ErrInSync
}

// hasChanged returns a boolean indicating whether a
// resource state has changed after processing
func (s *status) hasChanged(id string) bool {
	s.Lock()
	defer s.Unlock()

	return s.items[id] == nil
}

// hasFailed returns a boolean indicating whether a
// resource has failed during processing
func (s *status) hasFailed(id string) bool {
	s.Lock()
	defer s.Unlock()

	return s.items[id] != nil && s.items[id] != resource.ErrInSync
}

// New creates a new empty catalog with the provided configuration
func New(config *Config) *Catalog {
	c := &Catalog{
		config:     config,
		collection: make(resource.Collection),
		sorted:     make([]*graph.Node, 0),
		reversed:   graph.New(),
		status: &status{
			items: make(map[string]error),
		},
		Unsorted: make([]resource.Resource, 0),
	}

	// Inject the configuration for resources
	resource.DefaultConfig = &resource.Config{
		Logger:   config.Logger,
		SiteRepo: config.SiteRepo,
	}

	// Register the catalog type in Lua and also register
	// metamethods for the catalog, so that we can use
	// the catalog in a more Lua-friendly way
	mt := luar.MT(config.L, c)
	mt.RawSetString("__len", luar.New(config.L, (*Catalog).luaLen))
	config.L.SetGlobal("catalog", luar.New(config.L, c))

	return c
}

// Add adds a resource to the catalog.
// This method is called from Lua when adding new resources
func (c *Catalog) Add(resources ...resource.Resource) {
	for _, r := range resources {
		if r != nil {
			c.Unsorted = append(c.Unsorted, r)
		}
	}
}

// Load loads resources into the catalog
func (c *Catalog) Load() error {
	// Register the resource providers and catalog in Lua
	resource.LuaRegisterBuiltin(c.config.L)
	if err := c.config.L.DoFile(c.config.Module); err != nil {
		return err
	}

	// Perform a topological sort of the resources
	collection, err := resource.CreateCollection(c.Unsorted)
	if err != nil {
		return err
	}

	collectionGraph, err := collection.DependencyGraph()
	if err != nil {
		return err
	}

	reversed := collectionGraph.Reversed()

	sorted, err := collectionGraph.Sort()
	if err != nil {
		return err
	}

	// Set catalog fields
	c.collection = collection
	c.sorted = sorted
	c.reversed = reversed

	c.config.Logger.Printf("Loaded %d resources\n", len(c.sorted))

	return nil
}

// Run processes the resources from catalog
func (c *Catalog) Run() error {
	// process executes a single resource
	process := func(r resource.Resource) {
		id := r.ID()
		err := c.execute(r)
		c.status.set(id, err)
		if c.status.hasFailed(id) {
			c.config.Logger.Printf("%s %s\n", id, err)
			return
		}
		c.runTriggers(r)
	}

	// Start goroutines for concurrent processing
	var wg sync.WaitGroup
	ch := make(chan resource.Resource, 1024)
	c.config.Logger.Printf("Starting %d goroutines for concurrent processing\n", c.config.Concurrency)
	for i := 0; i < c.config.Concurrency; i++ {
		wg.Add(1)
		worker := func() {
			defer wg.Done()
			for r := range ch {
				c.config.Logger.Printf("%s is concurrent", r.ID())
				process(r)
			}
		}
		go worker()
	}

	// Process the resources
	for _, node := range c.sorted {
		r := c.collection[node.Name]
		switch {
		// Resource is concurrent and is an isolated node
		case r.IsConcurrent() && len(r.Dependencies()) == 0 && len(c.reversed.Nodes[r.ID()].Edges) == 0:
			ch <- r
			continue
		// Resource is concurrent and has no reverse dependencies
		case r.IsConcurrent() && len(c.reversed.Nodes[r.ID()].Edges) == 0:
			ch <- r
			continue
		// Resource is not concurrent
		default:
			process(r)
		}
	}

	close(ch)
	wg.Wait()

	// Print summary report
	if !c.config.DryRun {
		var changed, failed, uptodate int
		for _, err := range c.status.items {
			switch err {
			case nil:
				changed++
			case resource.ErrInSync:
				uptodate++
			default:
				failed++
			}
		}
		c.config.Logger.Printf("Resource summary is %d up-to-date, %d changed, %d failed\n", uptodate, changed, failed)
	}

	return nil
}

// execute processes a single resource
func (c *Catalog) execute(r resource.Resource) error {
	// Check if the resource has failed dependencies
	for _, dep := range r.Dependencies() {
		if c.status.hasFailed(dep) {
			return fmt.Errorf("failed dependency for %s", dep)
		}
	}

	if err := r.Validate(); err != nil {
		return err
	}

	state, err := r.Evaluate()
	if err != nil {
		return err
	}

	if c.config.DryRun {
		return nil
	}

	// Current and wanted states for the resource
	want := utils.NewString(state.Want)
	current := utils.NewString(state.Current)

	// The list of present and absent states for the resource
	present := utils.NewList(r.GetPresentStates()...)
	absent := utils.NewList(r.GetAbsentStates()...)

	id := r.ID()
	var action func() error
	switch {
	case want.IsInList(present) && current.IsInList(absent):
		action = r.Create
		c.config.Logger.Printf("%s is %s, should be %s\n", id, current, want)
	case want.IsInList(absent) && current.IsInList(present):
		action = r.Delete
		c.config.Logger.Printf("%s is %s, should be %s\n", id, current, want)
	case state.Outdated:
		action = r.Update
		c.config.Logger.Printf("%s is out of date\n", id)
	default:
		c.config.Logger.Printf("%s is up to date\n", id)
		return resource.ErrInSync
	}

	return action()
}

// runTriggers executes the triggers for each
// monitored resource if it's state has changed
func (c *Catalog) runTriggers(r resource.Resource) {
	for subscribed, trigger := range r.SubscribedTo() {
		if c.status.hasFailed(subscribed) {
			continue
		}

		c.config.Logger.Printf("%s running trigger, because %s has changed\n", r.ID(), subscribed)
		c.config.L.Push(trigger)
		if err := c.config.L.PCall(0, 0, nil); err != nil {
			c.config.Logger.Printf("%s trigger exited with an error: %s\n", r.ID(), err)
		}
	}
}

// luaLen returns the number of unsorted resources in catalog.
// This method is called from Lua.
func (c *Catalog) luaLen() int {
	return len(c.Unsorted)
}
<commit_msg>catalog: lock mutex while displaying summary of resources<commit_after>package catalog

import (
	"fmt"
	"log"
	"sync"

	"github.com/dnaeon/gru/graph"
	"github.com/dnaeon/gru/resource"
	"github.com/dnaeon/gru/utils"
	"github.com/layeh/gopher-luar"
	"github.com/yuin/gopher-lua"
)

// Catalog type contains a collection of resources
type Catalog struct {
	// Unsorted contains the list of resources created by Lua
	Unsorted []resource.Resource `luar:"-"`

	// Collection contains the unsorted resources as a collection
	collection resource.Collection `luar:"-"`

	// Sorted contains the resources after a topological sort.
	sorted []*graph.Node `luar:"-"`

	// Reversed contains the resource dependency graph in reverse
	// order. It is used for finding the reverse dependencies of
	// resources.
	reversed *graph.Graph `luar:"-"`

	// Status contains status information about resources
	status *status `luar:"-"`

	// Configuration settings
	config *Config `luar:"-"`
}

// Config type represents a set of settings to use when
// creating and processing the catalog
type Config struct {
	// Name of the Lua module to load and execute
	Module string

	// Do not take any actions, just report what would be done
	DryRun bool

	// Writer used to log events
	Logger *log.Logger

	// Path to the site repo containing module and data files
	SiteRepo string

	// The Lua state
	L *lua.LState

	// Number of goroutines to use for concurrent processing
	Concurrency int
}

// status type contains status information about processed resources
type status struct {
	sync.RWMutex

	// Items contain the result of resource processing and any
	// errors that might have occurred during processing.
	// Keys of the map are the resource ids and their
	// values are the errors returned by resources.
	items map[string]error
}

// set sets the status for a resource
func (s *status) set(id string, err error) {
	s.Lock()
	defer s.Unlock()
	s.items[id] = err
}

// get retrieves the status of a resource
func (s *status) get(id string) (error, bool) {
	s.Lock()
	defer s.Unlock()
	err, ok := s.items[id]

	return err, ok
}

// isSynced returns a boolean indicating whether a
// resource is up to date
func (s *status) isSynced(id string) bool {
	s.Lock()
	defer s.Unlock()

	return s.items[id] == resource.ErrInSync
}

// hasChanged returns a boolean indicating whether a
// resource state has changed after processing
func (s *status) hasChanged(id string) bool {
	s.Lock()
	defer s.Unlock()

	return s.items[id] == nil
}

// hasFailed returns a boolean indicating whether a
// resource has failed during processing
func (s *status) hasFailed(id string) bool {
	s.Lock()
	defer s.Unlock()

	return s.items[id] != nil && s.items[id] != resource.ErrInSync
}

// New creates a new empty catalog with the provided configuration
func New(config *Config) *Catalog {
	c := &Catalog{
		config:     config,
		collection: make(resource.Collection),
		sorted:     make([]*graph.Node, 0),
		reversed:   graph.New(),
		status: &status{
			items: make(map[string]error),
		},
		Unsorted: make([]resource.Resource, 0),
	}

	// Inject the configuration for resources
	resource.DefaultConfig = &resource.Config{
		Logger:   config.Logger,
		SiteRepo: config.SiteRepo,
	}

	// Register the catalog type in Lua and also register
	// metamethods for the catalog, so that we can use
	// the catalog in a more Lua-friendly way
	mt := luar.MT(config.L, c)
	mt.RawSetString("__len", luar.New(config.L, (*Catalog).luaLen))
	config.L.SetGlobal("catalog", luar.New(config.L, c))

	return c
}

// Add adds a resource to the catalog.
// This method is called from Lua when adding new resources
func (c *Catalog) Add(resources ...resource.Resource) {
	for _, r := range resources {
		if r != nil {
			c.Unsorted = append(c.Unsorted, r)
		}
	}
}

// Load loads resources into the catalog
func (c *Catalog) Load() error {
	// Register the resource providers and catalog in Lua
	resource.LuaRegisterBuiltin(c.config.L)
	if err := c.config.L.DoFile(c.config.Module); err != nil {
		return err
	}

	// Perform a topological sort of the resources
	collection, err := resource.CreateCollection(c.Unsorted)
	if err != nil {
		return err
	}

	collectionGraph, err := collection.DependencyGraph()
	if err != nil {
		return err
	}

	reversed := collectionGraph.Reversed()

	sorted, err := collectionGraph.Sort()
	if err != nil {
		return err
	}

	// Set catalog fields
	c.collection = collection
	c.sorted = sorted
	c.reversed = reversed

	c.config.Logger.Printf("Loaded %d resources\n", len(c.sorted))

	return nil
}

// Run processes the resources from catalog
func (c *Catalog) Run() error {
	// process executes a single resource
	process := func(r resource.Resource) {
		id := r.ID()
		err := c.execute(r)
		c.status.set(id, err)
		if c.status.hasFailed(id) {
			c.config.Logger.Printf("%s %s\n", id, err)
			return
		}
		c.runTriggers(r)
	}

	// Start goroutines for concurrent processing
	var wg sync.WaitGroup
	ch := make(chan resource.Resource, 1024)
	c.config.Logger.Printf("Starting %d goroutines for concurrent processing\n", c.config.Concurrency)
	for i := 0; i < c.config.Concurrency; i++ {
		wg.Add(1)
		worker := func() {
			defer wg.Done()
			for r := range ch {
				c.config.Logger.Printf("%s is concurrent", r.ID())
				process(r)
			}
		}
		go worker()
	}

	// Process the resources
	for _, node := range c.sorted {
		r := c.collection[node.Name]
		switch {
		// Resource is concurrent and is an isolated node
		case r.IsConcurrent() && len(r.Dependencies()) == 0 && len(c.reversed.Nodes[r.ID()].Edges) == 0:
			ch <- r
			continue
		// Resource is concurrent and has no reverse dependencies
		case r.IsConcurrent() && len(c.reversed.Nodes[r.ID()].Edges) == 0:
			ch <- r
			continue
		// Resource is not concurrent
		default:
			process(r)
		}
	}

	close(ch)
	wg.Wait()

	// Print summary report
	if !c.config.DryRun {
		c.status.Lock()
		var changed, failed, uptodate int
		for _, err := range c.status.items {
			switch err {
			case nil:
				changed++
			case resource.ErrInSync:
				uptodate++
			default:
				failed++
			}
		}
		c.config.Logger.Printf("Resource summary is %d up-to-date, %d changed, %d failed\n", uptodate, changed, failed)
		c.status.Unlock()
	}

	return nil
}

// execute processes a single resource
func (c *Catalog) execute(r resource.Resource) error {
	// Check if the resource has failed dependencies
	for _, dep := range r.Dependencies() {
		if c.status.hasFailed(dep) {
			return fmt.Errorf("failed dependency for %s", dep)
		}
	}

	if err := r.Validate(); err != nil {
		return err
	}

	state, err := r.Evaluate()
	if err != nil {
		return err
	}

	if c.config.DryRun {
		return nil
	}

	// Current and wanted states for the resource
	want := utils.NewString(state.Want)
	current := utils.NewString(state.Current)

	// The list of present and absent states for the resource
	present := utils.NewList(r.GetPresentStates()...)
	absent := utils.NewList(r.GetAbsentStates()...)

	id := r.ID()
	var action func() error
	switch {
	case want.IsInList(present) && current.IsInList(absent):
		action = r.Create
		c.config.Logger.Printf("%s is %s, should be %s\n", id, current, want)
	case want.IsInList(absent) && current.IsInList(present):
		action = r.Delete
		c.config.Logger.Printf("%s is %s, should be %s\n", id, current, want)
	case state.Outdated:
		action = r.Update
		c.config.Logger.Printf("%s is out of date\n", id)
	default:
		c.config.Logger.Printf("%s is up to date\n", id)
		return resource.ErrInSync
	}

	return action()
}

// runTriggers executes the triggers for each
// monitored resource if it's state has changed
func (c *Catalog) runTriggers(r resource.Resource) {
	for subscribed, trigger := range r.SubscribedTo() {
		if c.status.hasFailed(subscribed) {
			continue
		}

		c.config.Logger.Printf("%s running trigger, because %s has changed\n", r.ID(), subscribed)
		c.config.L.Push(trigger)
		if err := c.config.L.PCall(0, 0, nil); err != nil {
			c.config.Logger.Printf("%s trigger exited with an error: %s\n", r.ID(), err)
		}
	}
}

// luaLen returns the number of unsorted resources in catalog.
// This method is called from Lua.
func (c *Catalog) luaLen() int {
	return len(c.Unsorted)
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"io/ioutil"
	"testing"
)

type stubHTTPClient struct{}

func (stubHTTPClient) get(url string) ([]byte, error) {
	if url == "http://cacerts.digicert.com/DigiCertSHA2SecureServerCA.crt" {
		return ioutil.ReadFile("./testdata/issuer.pem")
	}

	return nil, errors.New("Unrecognised URL: " + url)
}

func TestGetIssuerCert(t *testing.T) {
	cert, err := readCertificate("./testdata/certificate.pem")
	if err != nil {
		t.Fatal(err)
	}

	client := &stubHTTPClient{}
	issCert, err := getIssuerCertificate(client, cert)
	if err != nil {
		t.Fatal(err)
	}

	if issCert.Issuer.CommonName != "DigiCert Global Root CA" {
		t.Fatal(issCert.Issuer.CommonName)
	}
}

func TestReadCertificate(t *testing.T) {
	_, err := readCertificate("./testdata/certificate.pem")
	if err != nil {
		t.Fatal(err)
	}
}

func TestGetOCSPServer(t *testing.T) {
	cert, _ := readCertificate("./testdata/certificate.pem")
	server, err := getOCSPServer(cert)
	if server != "http://ocsp.digicert.com" {
		t.Fatal(err)
	}
}

func TestDownloadCertificateUnreachable(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusServiceUnavailable)
	}))
	defer server.Close()

	_, err := downloadCertificate(server.URL)

	if err.Error() != "response code is not OK" {
		t.Fatalf("wrong error message: %s", err.Error())
	}
}

func TestCertificateFromBytesNoCertificate(t *testing.T) {
	in, _ := ioutil.ReadFile("./testdata/private_key.pem")
	_, err := certificateFromBytes(in)
	if err == nil {
		t.Fatal("should return error")
	}
}
<commit_msg>Removed unused test.<commit_after>package main

import (
	"errors"
	"io/ioutil"
	"testing"
)

type stubHTTPClient struct{}

func (stubHTTPClient) get(url string) ([]byte, error) {
	if url == "http://cacerts.digicert.com/DigiCertSHA2SecureServerCA.crt" {
		return ioutil.ReadFile("./testdata/issuer.pem")
	}

	return nil, errors.New("Unrecognised URL: " + url)
}

func TestGetIssuerCert(t *testing.T) {
	cert, err := readCertificate("./testdata/certificate.pem")
	if err != nil {
		t.Fatal(err)
	}

	client := &stubHTTPClient{}
	issCert, err := getIssuerCertificate(client, cert)
	if err != nil {
		t.Fatal(err)
	}

	if issCert.Issuer.CommonName != "DigiCert Global Root CA" {
		t.Fatal(issCert.Issuer.CommonName)
	}
}

func TestReadCertificate(t *testing.T) {
	_, err := readCertificate("./testdata/certificate.pem")
	if err != nil {
		t.Fatal(err)
	}
}

func TestGetOCSPServer(t *testing.T) {
	cert, _ := readCertificate("./testdata/certificate.pem")
	server, err := getOCSPServer(cert)
	if server != "http://ocsp.digicert.com" {
		t.Fatal(err)
	}
}

func TestCertificateFromBytesNoCertificate(t *testing.T) {
	in, _ := ioutil.ReadFile("./testdata/private_key.pem")
	_, err := certificateFromBytes(in)
	if err == nil {
		t.Fatal("should return error")
	}
}
<|endoftext|>


<commit_before>package index

import (
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
)

// Node represents a file tree.
//
// A single node represents a file or a directory.
//
// Nodes marked with EntryPromiseDel flag are marked
// as deleted, and are not going to be reachable via
// Lookup, Count methods. Deleting such nodes is a nop.
// This is how Node implements shallow delete.
type Node struct {
	Sub   map[string]*Node `json:"d"`
	Entry *Entry           `json:"e,omitempty"`
}

func newNode() *Node {
	return &Node{
		Sub:   make(map[string]*Node),
		Entry: NewEntry(0, 0755|os.ModeDir),
	}
}

// Add adds the given entry under the given path.
//
// Any deleted node, encountered on the tree path that, is going to
// be undeleted (having the EntryPromiseDel flag removed).
func (nd *Node) Add(path string, entry *Entry) {
	if path == "/" || path == "" {
		nd.Entry = entry
		return
	}

	var node string

	for {
		if nd.Deleted() {
			nd.undelete()
		}

		node, path = split(path)

		sub, ok := nd.Sub[node]
		if !ok {
			sub = newNode()
			nd.Sub[node] = sub
		}

		if path == "" {
			sub.Entry = entry
			return
		}

		nd = sub
	}
}

// Del disconnected a whole subtree rooted at a node given by the path.
//
// Del will ignore and do not disconnect nodes which are marked as deleted.
func (nd *Node) Del(path string) {
	var node string

	for {
		if nd.Deleted() {
			return
		}

		node, path = split(path)
		if path == "" {
			delete(nd.Sub, node)
			return
		}

		sub, ok := nd.Sub[node]
		if !ok {
			return
		}

		nd = sub
	}
}

// Clone creates a deep copy of node.
func (nd *Node) Clone() *Node {
	cpy := &Node{
		Sub: make(map[string]*Node, len(nd.Sub)),
	}

	for sub, node := range nd.Sub {
		if node != nil {
			cpy.Sub[sub] = node.Clone()
		} else {
			cpy.Sub[sub] = nil
		}
	}

	if nd.Entry != nil {
		cpy.Entry = nd.Entry.Clone()
	}

	return cpy
}

// PromiseAdd adds a node under the given path marked as newly added.
//
// If the node already exists, it'd be only marked with EntryPromiseAdd flag.
//
// If the node is already marked as newly added, the method is a no-op.
//
// If entry.Mode is non-zero, the effective node's entry is overwritten
// with this value.
//
// Rest of entry's fields are currently ignored.
func (nd *Node) PromiseAdd(path string, entry *Entry) {
	var newE *Entry

	if nd, ok := nd.lookup(path, true); ok {
		newE = nd.Entry
		newE.MergeIn(entry)
	} else {
		newE = NewEntry(entry.Size(), entry.Mode())
		newE.MergeIn(entry)
	}

	newE.SwapPromise(EntryPromiseAdd, EntryPromiseDel|EntryPromiseUnlink)
	nd.Add(path, newE)
}

// PromiseDel marks a node under the given path as deleted.
//
// If the node does not exist or is already marked as deleted, then
// method is no-op.
//
// If node is non-nil, then it's used instead of looking it up
// by the given path.
func (nd *Node) PromiseDel(path string, node *Node) {
	if node == nil {
		var ok bool
		node, ok = nd.Lookup(path)
		if !ok {
			return
		}
	}

	node.Entry.SwapPromise(EntryPromiseDel, EntryPromiseAdd|EntryPromiseUnlink)
}

// PromiseUnlink marks a node under the given path as unlinked.
//
// If the node does not exist or is already marked as unlinked, then
// method is no-op.
//
// If node is non-nil, then it's used instead of looking it up
// by the given path.
func (nd *Node) PromiseUnlink(path string, node *Node) {
	if node == nil {
		var ok bool
		node, ok = nd.Lookup(path)
		if !ok {
			return
		}
	}

	node.Entry.SwapPromise(EntryPromiseUnlink, EntryPromiseAdd)
}

// Count counts nodes which Entry.Size is at most maxsize.
//
// If maxsize is 0, the method is a no-op.
// If maxsize is < 0, the method counts all nodes.
//
// Count ignores nodes marked as deleted and/or virtual.
func (nd *Node) Count(maxsize int64) int {
	return nd.count(maxsize, false)
}

// CountAll counts nodes which Entry.Size is at most maxsize.
//
// If maxsize is 0, the method is a no-op.
// If maxsize is < 0, the method counts all nodes.
//
// CountAll does not ignore nodes marked as deleted and/or virtual.
func (nd *Node) CountAll(maxsize int64) int {
	return nd.count(maxsize, true)
}

func (nd *Node) count(maxsize int64, all bool) (count int) {
	if maxsize == 0 {
		return 0 // no-op
	}

	cur, stack := (*Node)(nil), []*Node{nd}

	for len(stack) != 0 {
		cur, stack = stack[0], stack[1:]

		if !all && (cur.Deleted() || cur.Virtual()) {
			continue
		}

		if cur.Entry != nil && (maxsize < 0 || cur.Entry.Size() <= maxsize) {
			count++
		}

		for _, nd := range cur.Sub {
			stack = append(stack, nd)
		}
	}

	return count
}

// DiskSize sums all Entry.Size of the nodes, given the condition the size
// is at most maxsize.
//
// If maxsize is 0, the method is a no-op.
// If maxsize is <0, all the nodes are summed up.
//
// DiskSize ignores nodes marked as deleted and/or virtual.
func (nd *Node) DiskSize(maxsize int64) (size int64) {
	return nd.diskSize(maxsize, false)
}

// DiskSizeAll sums all Entry.Size of the nodes, given the condition the size
// is at most maxsize.
//
// If maxsize is 0, the method is a no-op.
// If maxsize is <0, all the nodes are sumed up.
//
// DiskSizeAll does not ignore nodes marked as deleted and/or virtual.
func (nd *Node) DiskSizeAll(maxsize int64) (size int64) {
	return nd.diskSize(maxsize, true)
}

func (nd *Node) diskSize(maxsize int64, all bool) (size int64) {
	if maxsize == 0 {
		return 0 // no-op
	}

	stack := []*Node{nd}

	for len(stack) != 0 {
		nd, stack = stack[0], stack[1:]

		if !all && (nd.Deleted() || nd.Virtual()) {
			continue
		}

		if nd.Entry != nil && (maxsize < 0 || nd.Entry.Size() <= maxsize) {
			size += nd.Entry.Size()
		}

		for _, nd := range nd.Sub {
			stack = append(stack, nd)
		}
	}

	return size
}

// ForEach traverses the tree and calls fn on every node's entry.
//
// It ignored nodes marked as deleted.
func (nd *Node) ForEach(fn func(string, *Entry)) {
	nd.forEach(fn, false)
}

// ForEachAll traverses the tree and calls fn on every node's entry.
func (nd *Node) ForEachAll(fn func(string, *Entry)) {
	nd.forEach(fn, true)
}

func (nd *Node) forEach(fn func(string, *Entry), deleted bool) {
	type node struct {
		path string
		node *Node
	}

	n, stack := node{}, make([]node, 0, len(nd.Sub))

	// Add root node to stack.
	stack = append(stack, node{
		path: "",
		node: nd,
	})
	for len(stack) != 0 {
		n, stack = stack[0], stack[1:]

		if n.node.Deleted() && !deleted {
			continue
		}

		for path, nd := range n.node.Sub {
			stack = append(stack, node{
				path: filepath.Join(n.path, path),
				node: nd,
			})
		}

		fn(n.path, n.node.Entry)
	}
}

// Lookup looks up a node given by the path ignoring any of the node
// that is marked as deleted.
func (nd *Node) Lookup(path string) (*Node, bool) {
	return nd.lookup(path, false)
}

// LookupAll looks up a node given by the path.
func (nd *Node) LookupAll(path string) (*Node, bool) {
	return nd.lookup(path, true)
}

func (nd *Node) lookup(path string, deleted bool) (*Node, bool) {
	if path == "/" || path == "" {
		return nd.shallowCopy(), true
	}

	var node string

	for {
		if nd.Deleted() && !deleted {
			return nil, false
		}

		node, path = split(path)

		sub, ok := nd.Sub[node]
		if !ok {
			return nil, false
		}

		if path == "" {
			return sub.shallowCopy(), true
		}

		nd = sub
	}
}

// IsDir tells whether a node is a directory.
func (nd *Node) IsDir() bool {
	return nd.Entry.Mode().IsDir()
}

// Deleted tells whether node is marked as deleted.
func (nd *Node) Deleted() bool {
	return nd.Entry.Deleted()
}

// Virtual tells whether node is marked as virtual.
func (nd *Node) Virtual() bool {
	return nd.Entry.HasPromise(EntryPromiseVirtual)
}

func (nd *Node) undelete() {
	nd.Entry.SwapPromise(0, EntryPromiseDel|EntryPromiseUnlink)
}

func (nd *Node) shallowCopy() *Node {
	if nd.Sub != nil {
		sub := make(map[string]*Node, len(nd.Sub))

		for k, v := range nd.Sub {
			sub[k] = v
		}

		return &Node{
			Sub:   sub,
			Entry: nd.Entry,
		}
	}

	return nd
}

func split(path string) (string, string) {
	if path == "" {
		return "", ""
	}

	if path[0] == '/' {
		path = path[1:]
	}

	if i := strings.IndexRune(path, '/'); i != -1 {
		return path[:i], path[i+1:]
	}

	return path, ""
}

// UnmarshalJSON satisfies json.Unmarshaler interface. It initializes empty sub
// map when it's omitted in serialized data.
//
// Note: this is a fixing function that was created due to `omitempty` tag
// in Node's sub field.
func (nd *Node) UnmarshalJSON(data []byte) error {
	type tmp Node

	tnd := tmp{}
	if err := json.Unmarshal(data, &tnd); err != nil {
		return err
	}

	if *nd = Node(tnd); nd.Sub == nil {
		nd.Sub = make(map[string]*Node)
	}

	return nil
}
<commit_msg>klient/machine: static check for index.Node unmarshaller<commit_after>package index

import (
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
)

// Node represents a file tree.
//
// A single node represents a file or a directory.
//
// Nodes marked with EntryPromiseDel flag are marked
// as deleted, and are not going to be reachable via
// Lookup, Count methods. Deleting such nodes is a nop.
// This is how Node implements shallow delete.
type Node struct {
	Sub   map[string]*Node `json:"d"`
	Entry *Entry           `json:"e,omitempty"`
}

func newNode() *Node {
	return &Node{
		Sub:   make(map[string]*Node),
		Entry: NewEntry(0, 0755|os.ModeDir),
	}
}

// Add adds the given entry under the given path.
//
// Any deleted node, encountered on the tree path that, is going to
// be undeleted (having the EntryPromiseDel flag removed).
func (nd *Node) Add(path string, entry *Entry) {
	if path == "/" || path == "" {
		nd.Entry = entry
		return
	}

	var node string

	for {
		if nd.Deleted() {
			nd.undelete()
		}

		node, path = split(path)

		sub, ok := nd.Sub[node]
		if !ok {
			sub = newNode()
			nd.Sub[node] = sub
		}

		if path == "" {
			sub.Entry = entry
			return
		}

		nd = sub
	}
}

// Del disconnected a whole subtree rooted at a node given by the path.
//
// Del will ignore and do not disconnect nodes which are marked as deleted.
func (nd *Node) Del(path string) {
	var node string

	for {
		if nd.Deleted() {
			return
		}

		node, path = split(path)
		if path == "" {
			delete(nd.Sub, node)
			return
		}

		sub, ok := nd.Sub[node]
		if !ok {
			return
		}

		nd = sub
	}
}

// Clone creates a deep copy of node.
func (nd *Node) Clone() *Node {
	cpy := &Node{
		Sub: make(map[string]*Node, len(nd.Sub)),
	}

	for sub, node := range nd.Sub {
		if node != nil {
			cpy.Sub[sub] = node.Clone()
		} else {
			cpy.Sub[sub] = nil
		}
	}

	if nd.Entry != nil {
		cpy.Entry = nd.Entry.Clone()
	}

	return cpy
}

// PromiseAdd adds a node under the given path marked as newly added.
//
// If the node already exists, it'd be only marked with EntryPromiseAdd flag.
//
// If the node is already marked as newly added, the method is a no-op.
//
// If entry.Mode is non-zero, the effective node's entry is overwritten
// with this value.
//
// Rest of entry's fields are currently ignored.
func (nd *Node) PromiseAdd(path string, entry *Entry) {
	var newE *Entry

	if nd, ok := nd.lookup(path, true); ok {
		newE = nd.Entry
		newE.MergeIn(entry)
	} else {
		newE = NewEntry(entry.Size(), entry.Mode())
		newE.MergeIn(entry)
	}

	newE.SwapPromise(EntryPromiseAdd, EntryPromiseDel|EntryPromiseUnlink)
	nd.Add(path, newE)
}

// PromiseDel marks a node under the given path as deleted.
//
// If the node does not exist or is already marked as deleted, then
// method is no-op.
//
// If node is non-nil, then it's used instead of looking it up
// by the given path.
func (nd *Node) PromiseDel(path string, node *Node) {
	if node == nil {
		var ok bool
		node, ok = nd.Lookup(path)
		if !ok {
			return
		}
	}

	node.Entry.SwapPromise(EntryPromiseDel, EntryPromiseAdd|EntryPromiseUnlink)
}

// PromiseUnlink marks a node under the given path as unlinked.
//
// If the node does not exist or is already marked as unlinked, then
// method is no-op.
//
// If node is non-nil, then it's used instead of looking it up
// by the given path.
func (nd *Node) PromiseUnlink(path string, node *Node) {
	if node == nil {
		var ok bool
		node, ok = nd.Lookup(path)
		if !ok {
			return
		}
	}

	node.Entry.SwapPromise(EntryPromiseUnlink, EntryPromiseAdd)
}

// Count counts nodes which Entry.Size is at most maxsize.
//
// If maxsize is 0, the method is a no-op.
// If maxsize is < 0, the method counts all nodes.
//
// Count ignores nodes marked as deleted and/or virtual.
func (nd *Node) Count(maxsize int64) int {
	return nd.count(maxsize, false)
}

// CountAll counts nodes which Entry.Size is at most maxsize.
//
// If maxsize is 0, the method is a no-op.
// If maxsize is < 0, the method counts all nodes.
//
// CountAll does not ignore nodes marked as deleted and/or virtual.
func (nd *Node) CountAll(maxsize int64) int {
	return nd.count(maxsize, true)
}

func (nd *Node) count(maxsize int64, all bool) (count int) {
	if maxsize == 0 {
		return 0 // no-op
	}

	cur, stack := (*Node)(nil), []*Node{nd}

	for len(stack) != 0 {
		cur, stack = stack[0], stack[1:]

		if !all && (cur.Deleted() || cur.Virtual()) {
			continue
		}

		if cur.Entry != nil && (maxsize < 0 || cur.Entry.Size() <= maxsize) {
			count++
		}

		for _, nd := range cur.Sub {
			stack = append(stack, nd)
		}
	}

	return count
}

// DiskSize sums all Entry.Size of the nodes, given the condition the size
// is at most maxsize.
//
// If maxsize is 0, the method is a no-op.
// If maxsize is <0, all the nodes are summed up.
//
// DiskSize ignores nodes marked as deleted and/or virtual.
func (nd *Node) DiskSize(maxsize int64) (size int64) {
	return nd.diskSize(maxsize, false)
}

// DiskSizeAll sums all Entry.Size of the nodes, given the condition the size
// is at most maxsize.
//
// If maxsize is 0, the method is a no-op.
// If maxsize is <0, all the nodes are sumed up.
//
// DiskSizeAll does not ignore nodes marked as deleted and/or virtual.
func (nd *Node) DiskSizeAll(maxsize int64) (size int64) {
	return nd.diskSize(maxsize, true)
}

func (nd *Node) diskSize(maxsize int64, all bool) (size int64) {
	if maxsize == 0 {
		return 0 // no-op
	}

	stack := []*Node{nd}

	for len(stack) != 0 {
		nd, stack = stack[0], stack[1:]

		if !all && (nd.Deleted() || nd.Virtual()) {
			continue
		}

		if nd.Entry != nil && (maxsize < 0 || nd.Entry.Size() <= maxsize) {
			size += nd.Entry.Size()
		}

		for _, nd := range nd.Sub {
			stack = append(stack, nd)
		}
	}

	return size
}

// ForEach traverses the tree and calls fn on every node's entry.
//
// It ignored nodes marked as deleted.
func (nd *Node) ForEach(fn func(string, *Entry)) {
	nd.forEach(fn, false)
}

// ForEachAll traverses the tree and calls fn on every node's entry.
func (nd *Node) ForEachAll(fn func(string, *Entry)) {
	nd.forEach(fn, true)
}

func (nd *Node) forEach(fn func(string, *Entry), deleted bool) {
	type node struct {
		path string
		node *Node
	}

	n, stack := node{}, make([]node, 0, len(nd.Sub))

	// Add root node to stack.
	stack = append(stack, node{
		path: "",
		node: nd,
	})
	for len(stack) != 0 {
		n, stack = stack[0], stack[1:]

		if n.node.Deleted() && !deleted {
			continue
		}

		for path, nd := range n.node.Sub {
			stack = append(stack, node{
				path: filepath.Join(n.path, path),
				node: nd,
			})
		}

		fn(n.path, n.node.Entry)
	}
}

// Lookup looks up a node given by the path ignoring any of the node
// that is marked as deleted.
func (nd *Node) Lookup(path string) (*Node, bool) {
	return nd.lookup(path, false)
}

// LookupAll looks up a node given by the path.
func (nd *Node) LookupAll(path string) (*Node, bool) {
	return nd.lookup(path, true)
}

func (nd *Node) lookup(path string, deleted bool) (*Node, bool) {
	if path == "/" || path == "" {
		return nd.shallowCopy(), true
	}

	var node string

	for {
		if nd.Deleted() && !deleted {
			return nil, false
		}

		node, path = split(path)

		sub, ok := nd.Sub[node]
		if !ok {
			return nil, false
		}

		if path == "" {
			return sub.shallowCopy(), true
		}

		nd = sub
	}
}

// IsDir tells whether a node is a directory.
func (nd *Node) IsDir() bool {
	return nd.Entry.Mode().IsDir()
}

// Deleted tells whether node is marked as deleted.
func (nd *Node) Deleted() bool {
	return nd.Entry.Deleted()
}

// Virtual tells whether node is marked as virtual.
func (nd *Node) Virtual() bool {
	return nd.Entry.HasPromise(EntryPromiseVirtual)
}

func (nd *Node) undelete() {
	nd.Entry.SwapPromise(0, EntryPromiseDel|EntryPromiseUnlink)
}

func (nd *Node) shallowCopy() *Node {
	if nd.Sub != nil {
		sub := make(map[string]*Node, len(nd.Sub))

		for k, v := range nd.Sub {
			sub[k] = v
		}

		return &Node{
			Sub:   sub,
			Entry: nd.Entry,
		}
	}

	return nd
}

func split(path string) (string, string) {
	if path == "" {
		return "", ""
	}

	if path[0] == '/' {
		path = path[1:]
	}

	if i := strings.IndexRune(path, '/'); i != -1 {
		return path[:i], path[i+1:]
	}

	return path, ""
}

var _ json.Unmarshaler = (*Node)(nil)

// UnmarshalJSON satisfies json.Unmarshaler interface. It initializes empty sub
// map when it's omitted in serialized data.
//
// Note: this is a fixing function that was created due to `omitempty` tag
// in Node's sub field.
func (nd *Node) UnmarshalJSON(data []byte) error {
	type tmp Node

	tnd := tmp{}
	if err := json.Unmarshal(data, &tnd); err != nil {
		return err
	}

	if *nd = Node(tnd); nd.Sub == nil {
		nd.Sub = make(map[string]*Node)
	}

	return nil
}
<|endoftext|>


<commit_before>// ############## Page 169-172 ##############
package main

import "fmt"

/*
  Bicycle
  Automatic delegation to Spares() is provided by Parts.
*/
type Bicycle struct {
  Size string
  Parts
}

/*
  Parts
  Parts is a slice vs. a struct containing one, giving it slice like behavior.
  (like on Page 173)
*/
type Parts []Part

func (parts Parts) Spares() (spares []Part) {
  for _, part := range parts {
    if part.NeedsSpare {
      spares = append(spares, part)
    }
  }
  return spares
}

/*
  Part
  NeedsSpare defaults to false here, the opposite of the Ruby example. We could
  write a NewSparePart initializer that defaults it to true.
*/
type Part struct {
  Name        string
  Description string
  NeedsSpare  bool
}

func main() {
  chain := Part{Name: "chain", Description: "10-speed", NeedsSpare: true}
  roadTire := Part{Name: "tire_size", Description: "23", NeedsSpare: true}
  tape := Part{Name: "tape_color", Description: "red", NeedsSpare: true}
  mountainTire := Part{Name: "tire_size", Description: "2.1", NeedsSpare: true}
  rearShock := Part{Name: "rear_shock", Description: "Fox", NeedsSpare: true}
  frontShock := Part{Name: "front_shock", Description: "Manitou"}

  roadBikeParts := Parts{chain, roadTire, tape}
  roadBike := Bicycle{Size: "L", Parts: roadBikeParts}

  mountainBikeParts := Parts{chain, mountainTire, frontShock, rearShock}
  mountainBike := Bicycle{Size: "L", Parts: mountainBikeParts}

  fmt.Println(roadBike.Size)
  fmt.Println(roadBike.Spares())

  fmt.Println(mountainBike.Size)
  fmt.Println(mountainBike.Spares())
}
<commit_msg>combining Parts<commit_after>// ############## Page 169-172 ##############
package main

import "fmt"

/*
  Bicycle
  Automatic delegation to Spares() is provided by Parts.
*/
type Bicycle struct {
  Size string
  Parts
}

/*
  Parts
  Parts is a slice vs. a struct containing one, giving it slice like behavior.
  (like on Page 173)
*/
type Parts []Part

func (parts Parts) Spares() (spares []Part) {
  for _, part := range parts {
    if part.NeedsSpare {
      spares = append(spares, part)
    }
  }
  return spares
}

/*
  Part
  NeedsSpare defaults to false here, the opposite of the Ruby example. We could
  write a NewSparePart initializer that defaults it to true.
*/
type Part struct {
  Name        string
  Description string
  NeedsSpare  bool
}

func main() {
  chain := Part{Name: "chain", Description: "10-speed", NeedsSpare: true}
  roadTire := Part{Name: "tire_size", Description: "23", NeedsSpare: true}
  tape := Part{Name: "tape_color", Description: "red", NeedsSpare: true}
  mountainTire := Part{Name: "tire_size", Description: "2.1", NeedsSpare: true}
  rearShock := Part{Name: "rear_shock", Description: "Fox", NeedsSpare: true}
  frontShock := Part{Name: "front_shock", Description: "Manitou"}

  roadBikeParts := Parts{chain, roadTire, tape}
  roadBike := Bicycle{Size: "L", Parts: roadBikeParts}

  mountainBikeParts := Parts{chain, mountainTire, frontShock, rearShock}
  mountainBike := Bicycle{Size: "L", Parts: mountainBikeParts}

  fmt.Println(roadBike.Size)
  fmt.Println(roadBike.Spares())

  fmt.Println(mountainBike.Size)
  fmt.Println(mountainBike.Spares())

  // We can combine Parts and still call Spares, unlike the Ruby example. (page 173)
  comboParts := Parts{}
  comboParts = append(comboParts, mountainBike.Parts...)
  comboParts = append(comboParts, roadBike.Parts...)
  fmt.Println(len(comboParts))
  fmt.Println(comboParts.Spares())
}
<|endoftext|>


<commit_before>package checker

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/asciimoo/privacyscore/result"
	"github.com/asciimoo/privacyscore/utils"
)

const RESOURCE_LIMIT = 64
const TIMEOUT = 5
const USER_AGENT = "Mozilla/5.0 (compatible) PrivacyScore Checker v0.1.0"
const maxResponseBodySize = 1024 * 1024 * 5

var mutex = &sync.Mutex{}

type Checker interface {
	Check(*CheckJob, *PageInfo)
}

type CheckJob struct {
	sync.RWMutex
	Result    *result.Result
	Resources map[string]*PageInfo
	Chan      chan bool
}

type PageInfo struct {
	ResponseBody   []byte
	ContentType    string
	StatusCode     int
	URL            *url.URL
	OriginalURL    *url.URL
	Cookies        []*http.Cookie
	Domain         string
	ResponseHeader *http.Header
}

var checkers []Checker = []Checker{
	&CookieChecker{},
	&HTMLChecker{},
	&HTTPSChecker{},
	&SecureHeaderChecker{},
	&CSSChecker{},
}

func Run(URL string) (*CheckJob, error) {
	if !strings.HasPrefix(URL, "http://") && !strings.HasPrefix(URL, "https://") {
		URL = "http://" + URL
	}
	c := newCheckJob(URL)
	finishedResources := 0
	errorCount := 0
	c.CheckURL(URL)
	for finishedResources != len(c.Resources) && finishedResources < RESOURCE_LIMIT {
		select {
		case ret := <-c.Chan:
			if ret == false {
				errorCount += 1
			}
			finishedResources += 1
		}
	}
	if finishedResources == 0 || (errorCount > 0 && errorCount == finishedResources) {
		return c, errors.New("Could not download host")
	}
	c.RLock()
	r, found := c.Resources[URL]
	c.RUnlock()
	if found && r != nil {
		c.Result.BaseURL = r.URL.String()
	}
	return c, nil
}

func newCheckJob(URL string) *CheckJob {
	return &CheckJob{
		Result:    result.New(URL),
		Resources: make(map[string]*PageInfo),
		Chan:      make(chan bool, RESOURCE_LIMIT),
	}
}

func (c *CheckJob) CheckURL(URL string) {
	if URL == "" {
		return
	}
	// URL already added
	c.RLock()
	r, found := c.Resources[URL]
	c.RUnlock()
	if found && r != nil {
		return
	}
	// resource limit exceeded
	if len(c.Resources) >= RESOURCE_LIMIT {
		c.Result.AddError(errors.New(fmt.Sprintf("Resource limit exceeded, site loads more than %v resources", RESOURCE_LIMIT)))
		return
	}
	var p *PageInfo
	c.Lock()
	c.Resources[URL] = p
	c.Unlock()
	go func() {
		r, err := fetchURL(URL)
		if err != nil {
			c.Result.AddError(err)
			c.Chan <- false
			return
		}
		var body []byte
		contentType := r.Header.Get("Content-Type")
		if r.StatusCode == 200 && (strings.Contains(contentType, "text/") || strings.Contains(contentType, "html")) {
			body, err = ioutil.ReadAll(io.LimitReader(r.Body, maxResponseBodySize))
			if err != nil {
				c.Result.AddError(err)
			}
		} else {
			body = []byte{}
		}
		r.Body.Close()

		u, _ := url.Parse(URL)
		p = &PageInfo{
			body,
			r.Header.Get("Content-Type"),
			r.StatusCode,
			r.Request.URL,
			u,
			r.Cookies(),
			utils.CropSubdomains(r.Request.URL.Host),
			&r.Header,
		}
		c.Lock()
		c.Resources[URL] = p
		c.Unlock()
		for _, ch := range checkers {
			ch.Check(c, p)
		}
		c.Chan <- true
	}()
	return
}

func fetchURL(URL string) (*http.Response, error) {
	client := http.Client{
		Timeout: time.Duration(TIMEOUT * time.Second),
		Transport: &http.Transport{
			DisableKeepAlives:     true,
			ResponseHeaderTimeout: TIMEOUT * time.Second,
		},
	}
	req, err := http.NewRequest("GET", URL, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", USER_AGENT)
	response, err := client.Do(req)
	return response, err
}
<commit_msg>[fix] properly detect the end of a check<commit_after>package checker

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/asciimoo/privacyscore/result"
	"github.com/asciimoo/privacyscore/utils"
)

const RESOURCE_LIMIT = 64
const TIMEOUT = 5
const USER_AGENT = "Mozilla/5.0 (compatible) PrivacyScore Checker v0.1.0"
const maxResponseBodySize = 1024 * 1024 * 5

var mutex = &sync.Mutex{}

type Checker interface {
	Check(*CheckJob, *PageInfo)
}

type CheckJob struct {
	sync.RWMutex
	Result    *result.Result
	Resources map[string]*PageInfo
	wg        *sync.WaitGroup
}

type PageInfo struct {
	ResponseBody   []byte
	ContentType    string
	StatusCode     int
	URL            *url.URL
	OriginalURL    *url.URL
	Cookies        []*http.Cookie
	Domain         string
	ResponseHeader *http.Header
}

var checkers []Checker = []Checker{
	&CookieChecker{},
	&HTMLChecker{},
	&HTTPSChecker{},
	&SecureHeaderChecker{},
	&CSSChecker{},
}

func Run(URL string) (*CheckJob, error) {
	if !strings.HasPrefix(URL, "http://") && !strings.HasPrefix(URL, "https://") {
		URL = "http://" + URL
	}
	c := newCheckJob(URL)
	r, found := c.Resources[URL]
	if found && r != nil {
		c.Result.BaseURL = r.URL.String()
		return c, nil
	}
	return c, errors.New("Could not download host")
}

func newCheckJob(URL string) *CheckJob {
	c := &CheckJob{
		Result:    result.New(URL),
		Resources: make(map[string]*PageInfo),
		wg:        &sync.WaitGroup{},
	}
	c.Check(URL)
	return c
}

func (c *CheckJob) Check(URL string) {
	c.CheckURL(URL)
	c.wg.Wait()
}

func (c *CheckJob) CheckURL(URL string) {
	if URL == "" {
		return
	}
	// URL already added
	c.RLock()
	r, found := c.Resources[URL]
	c.RUnlock()
	if found && r != nil {
		return
	}
	// resource limit exceeded
	if len(c.Resources) >= RESOURCE_LIMIT {
		c.Result.AddError(errors.New(fmt.Sprintf("Resource limit exceeded, site loads more than %v resources", RESOURCE_LIMIT)))
		return
	}
	var p *PageInfo
	c.Lock()
	c.Resources[URL] = p
	c.wg.Add(1)
	c.Unlock()
	go func() {
		r, err := fetchURL(URL)
		if err != nil {
			c.Result.AddError(err)
			c.wg.Done()
			return
		}
		var body []byte
		contentType := r.Header.Get("Content-Type")
		if r.StatusCode == 200 && (strings.Contains(contentType, "text/") || strings.Contains(contentType, "html")) {
			body, err = ioutil.ReadAll(io.LimitReader(r.Body, maxResponseBodySize))
			if err != nil {
				c.Result.AddError(err)
			}
		} else {
			body = []byte{}
		}
		r.Body.Close()

		u, _ := url.Parse(URL)
		p = &PageInfo{
			body,
			r.Header.Get("Content-Type"),
			r.StatusCode,
			r.Request.URL,
			u,
			r.Cookies(),
			utils.CropSubdomains(r.Request.URL.Host),
			&r.Header,
		}
		c.Lock()
		c.Resources[URL] = p
		c.Unlock()
		for _, ch := range checkers {
			ch.Check(c, p)
		}
		c.wg.Done()
	}()
	return
}

func fetchURL(URL string) (*http.Response, error) {
	client := http.Client{
		Timeout: time.Duration(TIMEOUT * time.Second),
		Transport: &http.Transport{
			DisableKeepAlives:     true,
			ResponseHeaderTimeout: TIMEOUT * time.Second,
		},
	}
	req, err := http.NewRequest("GET", URL, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", USER_AGENT)
	response, err := client.Do(req)
	return response, err
}
<|endoftext|>


<commit_before>package checker

import (
	"errors"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/asciimoo/privacyscore/result"
	"github.com/asciimoo/privacyscore/utils"
)

const USER_AGENT = "Mozilla/5.0 (compatible) PrivacyScore Checker v0.1.0"
const TIMEOUT = 5
const maxResponseBodySize = 1024 * 1024 * 5

var mutex = &sync.Mutex{}

type Checker interface {
	Check(*result.Result, *PageInfo)
}

type Check struct {
	sync.RWMutex
	Result    *result.Result
	Resources map[string]*PageInfo
}

type PageInfo struct {
	ResponseBody   []byte
	ContentType    string
	StatusCode     int
	URL            *url.URL
	OriginalURL    *url.URL
	Cookies        []*http.Cookie
	Domain         string
	ResponseHeader *http.Header
}

var checkers []Checker = []Checker{
	&CookieChecker{},
	&HTMLChecker{},
	&HTTPSChecker{},
	&SecureHeaderChecker{},
}

func Run(URL string) (*result.Result, error) {
	if !strings.HasPrefix(URL, "http://") && !strings.HasPrefix(URL, "https://") {
		URL = "http://" + URL
	}
	c := newCheck(URL)
	_, err := c.CheckURL(URL)
	return c.Result, err
}

func newCheck(URL string) *Check {
	return &Check{
		Result:    result.New(URL),
		Resources: make(map[string]*PageInfo),
	}
}

func (c *Check) CheckURL(URL string) (*PageInfo, error) {
	if u, found := c.Resources[URL]; found {
		return u, errors.New("URL already added")
	}
	r, err := fetchURL(URL)
	if err != nil {
		return nil, err
	}
	defer r.Body.Close()
	body, err := ioutil.ReadAll(io.LimitReader(r.Body, maxResponseBodySize))
	u, _ := url.Parse(URL)
	p := &PageInfo{
		body,
		r.Header.Get("Content-Type"),
		r.StatusCode,
		r.Request.URL,
		u,
		r.Cookies(),
		utils.CropSubdomains(r.Request.URL.Host),
		&r.Header,
	}
	if err != nil {
		c.Result.AddError(err)
	}
	c.Lock()
	c.Resources[URL] = p
	c.Unlock()
	for _, ch := range checkers {
		ch.Check(c.Result, p)
	}
	return p, nil
}

func fetchURL(URL string) (*http.Response, error) {
	client := http.Client{
		Timeout: time.Duration(TIMEOUT * time.Second),
		Transport: &http.Transport{
			DisableKeepAlives: true,
		},
	}
	req, err := http.NewRequest("GET", URL, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", USER_AGENT)
	response, err := client.Do(req)
	return response, err
}
<commit_msg>[enh] async checks<commit_after>package checker

import (
	"errors"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/asciimoo/privacyscore/result"
	"github.com/asciimoo/privacyscore/utils"
)

const RESOURCE_LIMIT = 64
const TIMEOUT = 5
const USER_AGENT = "Mozilla/5.0 (compatible) PrivacyScore Checker v0.1.0"
const maxResponseBodySize = 1024 * 1024 * 5

var mutex = &sync.Mutex{}

type Checker interface {
	Check(*result.Result, *PageInfo)
}

type CheckJob struct {
	sync.RWMutex
	Result    *result.Result
	Resources map[string]*PageInfo
	Chan      chan bool
}

type PageInfo struct {
	ResponseBody   []byte
	ContentType    string
	StatusCode     int
	URL            *url.URL
	OriginalURL    *url.URL
	Cookies        []*http.Cookie
	Domain         string
	ResponseHeader *http.Header
}

var checkers []Checker = []Checker{
	&CookieChecker{},
	&HTMLChecker{},
	&HTTPSChecker{},
	&SecureHeaderChecker{},
	&CSSChecker{},
}

func Run(URL string) (*result.Result, error) {
	if !strings.HasPrefix(URL, "http://") && !strings.HasPrefix(URL, "https://") {
		URL = "http://" + URL
	}
	c := newCheckJob(URL)
	finishedResources := 0
	c.CheckURL(URL)
	for finishedResources != len(c.Resources) && finishedResources < RESOURCE_LIMIT {
		select {
		case _ = <-c.Chan:
			finishedResources += 1
		}
	}
	if finishedResources == 0 {
		return c.Result, errors.New("Could not download host")
	}
	return c.Result, nil
}

func newCheckJob(URL string) *CheckJob {
	return &CheckJob{
		Result:    result.New(URL),
		Resources: make(map[string]*PageInfo),
		Chan:      make(chan bool, RESOURCE_LIMIT),
	}
}

func (c *CheckJob) CheckURL(URL string) {
	// URL already added
	if _, found := c.Resources[URL]; found {
		c.Chan <- true
		return
	}
	r, err := fetchURL(URL)
	if err != nil {
		c.Result.AddError(err)
		c.Chan <- false
		return
	}
	defer r.Body.Close()
	body, err := ioutil.ReadAll(io.LimitReader(r.Body, maxResponseBodySize))
	u, _ := url.Parse(URL)
	p := &PageInfo{
		body,
		r.Header.Get("Content-Type"),
		r.StatusCode,
		r.Request.URL,
		u,
		r.Cookies(),
		utils.CropSubdomains(r.Request.URL.Host),
		&r.Header,
	}
	if err != nil {
		c.Result.AddError(err)
	}
	c.Lock()
	c.Resources[URL] = p
	c.Unlock()
	go func() {
		for _, ch := range checkers {
			ch.Check(c.Result, p)
		}
		c.Chan <- false
	}()
	return
}

func fetchURL(URL string) (*http.Response, error) {
	client := http.Client{
		Timeout: time.Duration(TIMEOUT * time.Second),
		Transport: &http.Transport{
			DisableKeepAlives: true,
		},
	}
	req, err := http.NewRequest("GET", URL, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", USER_AGENT)
	response, err := client.Do(req)
	return response, err
}
<|endoftext|>


<commit_before>/*
 * The Clammit application intercepts HTTP POST requests with content-type
 * "multipart/form-data", forwards any "file" form-data elements to ClamAV
 * and only forwards the request to the application if ClamAV passes all
 * of these elements as virus-free.
 */
package main

import (
	"bytes"
	"clammit/forwarder"
	"clammit/scanner"
	"encoding/json"
	"flag"
	"fmt"
	"gopkg.in/gcfg.v1"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"runtime"
	"strconv"
	"strings"
	"syscall"
)

//
// Configuration structure, designed for gcfg
//
type Config struct {
	App ApplicationConfig `gcfg:"application"`
}

type ApplicationConfig struct {
	// The address to listen on. This can be one of:
	// * tcp:host:port
	// * tcp:port
	// * unix:filename
	// * host:port
	// * :port
	//
	// For example:
	//   Listen: tcp:0.0.0.0:8438
	//   Listen: unix:/tmp/clammit.sock
	//   Listen: :8438
	Listen string `gcfg:"listen"`
	// Socket file permissions (only used if listening on a unix socket), in octal form.
	//
	// For example:
	//   SocketPerms: 0766
	SocketPerms string `gcfg:"unix-socket-perms"`
	// The URL of the application that Clammit is proxying. Generally, this will
	// be the base URL (http://host:port/), but you can also add a path prefix
	// if needed (http://host:port/prefix)
	ApplicationURL string `gcfg:"application-url"`
	// The URL of clamd, which will either be TCP or Unix:
	//
	// For example:
	//   ClamdURL: tcp://localhost:3310
	//   ClamdURL: unix:/tmp/clamd.sock
	ClamdURL string `gcfg:"clamd-url"`
	// The HTTP status code to return when a virus is found
	VirusStatusCode int `gcfg:"virus-status-code"`
	// If the body content-length exceeds this value, it will be written to
	// disk. Below it, we'll hold the whole body in memory to improve speed.
	ContentMemoryThreshold int64 `gcfg:"content-memory-threshold"`
	// Log file name (default is to log to stdout)
	Logfile string `gcfg:"log-file"`
	// If true, clammit will expose a small test HTML page.
	TestPages bool `gcfg:"test-pages"`
	// If true, will log the progression of each request through the forwarder
	Debug bool `gcfg:"debug"`
	// If true, will log the annoying clamd messages
	DebugClam bool `gcfg:"debug-clam"`
	// Number of CPU threads to use
	NumThreads int `gcfg:"num-threads"`
}

//
// Default configuration
//
var DefaultApplicationConfig = ApplicationConfig{
	Listen:                 ":8438",
	SocketPerms:            "0777",
	ApplicationURL:         "",
	ClamdURL:               "",
	VirusStatusCode:        418,
	ContentMemoryThreshold: 1024 * 1024,
	Logfile:                "",
	TestPages:              true,
	Debug:                  false,
	DebugClam:              false,
	NumThreads:             runtime.NumCPU(),
}

//
// Application context
//
type Ctx struct {
	Config          Config
	ApplicationURL  *url.URL
	ScanInterceptor *ScanInterceptor
	Scanner         scanner.Scanner
	Logger          *log.Logger
	Listener        net.Listener
	ActivityChan    chan int
	ShuttingDown    bool
}

//
// JSON server information response
//
type Info struct {
	Address             string `json:"scan_server_url"`
	PingResult          string `json:"ping_result"`
	Version             string `json:"version"`
	TestScanVirusResult string `json:"test_scan_virus"`
	TestScanCleanResult string `json:"test_scan_clean"`
}

//
// Global variables and config
//
var ctx *Ctx
var configFile string
var EICAR = []byte(`X5O!P%@AP[4\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*`)

func init() {
	flag.StringVar(&configFile, "config", "", "Configuration file")
}

func main() {
	/*
	 * Construct configuration, set up logging
	 */
	flag.Parse()
	ctx = &Ctx{
		ActivityChan: make(chan int),
		ShuttingDown: false,
	}

	ctx.Config.App = DefaultApplicationConfig

	if err := gcfg.ReadFileInto(&ctx.Config, configFile); err != nil {
		log.Fatalf("Configuration read failure: %s", err.Error())
	}

	// Socket perms are octal!
	socketPerms := 0777
	if ctx.Config.App.SocketPerms != "" {
		if sp, err := strconv.ParseInt(ctx.Config.App.SocketPerms, 8, 0); err == nil {
			socketPerms = int(sp)
		} else {
			log.Fatalf("SocketPerms invalid (expected 4-digit octal: %s", err.Error)
		}
	}

	// Allow multi-proc
	runtime.GOMAXPROCS(ctx.Config.App.NumThreads)

	startLogging()

	/*
	 * Construct objects, validate the URLs
	 */
	ctx.ApplicationURL = checkURL(ctx.Config.App.ApplicationURL)
	checkURL(ctx.Config.App.ClamdURL)

	ctx.Scanner = new(scanner.Clamav)
	ctx.Scanner.SetLogger(ctx.Logger, ctx.Config.App.Debug)
	ctx.Scanner.SetAddress(ctx.Config.App.ClamdURL)

	ctx.ScanInterceptor = &ScanInterceptor{
		VirusStatusCode: ctx.Config.App.VirusStatusCode,
		Scanner:         ctx.Scanner,
	}

	/*
	 * Set up the HTTP server
	 */
	router := http.NewServeMux()

	router.HandleFunc("/clammit", infoHandler)
	router.HandleFunc("/clammit/scan", scanHandler)
	if ctx.Config.App.TestPages {
		fs := http.FileServer(http.Dir("testfiles"))
		router.Handle("/clammit/test/", http.StripPrefix("/clammit/test/", fs))
	}
	router.HandleFunc("/", scanForwardHandler)

	if !ctx.Config.App.DebugClam {
		log.SetOutput(ioutil.Discard) // go-clamd has irritating logging, so turn it off
	}

	if listener, err := getListener(ctx.Config.App.Listen, socketPerms); err != nil {
		ctx.Logger.Fatal("Unable to listen on: ", ctx.Config.App.Listen, ", reason: ", err)
	} else {
		ctx.Listener = listener
		beGraceful() // graceful shutdown from here on in
		ctx.Logger.Println("Listening on", ctx.Config.App.Listen)
		http.Serve(listener, router)
	}
}

/*
 * Starts logging
 */
func startLogging() {
	if ctx.Config.App.Logfile != "" {
		w, err := os.OpenFile(ctx.Config.App.Logfile, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0660)
		if err == nil {
			ctx.Logger = log.New(w, "", log.LstdFlags)
		} else {
			log.Fatal("Failed to open log file", ctx.Config.App.Logfile, ":", err)
		}
	} else {
		ctx.Logger = log.New(os.Stdout, "", log.LstdFlags)
		ctx.Logger.Println("No log file configured - using stdout")
	}
}

/*
 * Handles graceful shutdown. Sets ctx.ShuttingDown = true to stop any new
 * requests, then waits for active requests to complete before closing the
 * HTTP listener.
 */
func beGraceful() {
	sigchan := make(chan os.Signal)
	signal.Notify(sigchan, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		activity := 0
		for {
			select {
			case _ = <-sigchan:
				ctx.Logger.Println("Received termination signal")
				ctx.ShuttingDown = true
				for activity > 0 {
					ctx.Logger.Printf("There are %d active requests, waiting", activity)
					i := <-ctx.ActivityChan
					activity += i
				}
				// This will cause main() to continue from http.Serve()
				// it will also clean up the unix socket (if relevant)
				ctx.Listener.Close()
			case i := <-ctx.ActivityChan:
				activity += i
			}
		}
	}()
}

/*
 * Validates the URL is OK (fatal error if not) and returns it
 */
func checkURL(urlString string) *url.URL {
	parsedURL, err := url.Parse(urlString)
	if err != nil {
		log.Fatal("Invalid URL:", urlString)
	}
	return parsedURL
}

/*
 * Returns a TCP or Unix socket listener, according to the scheme prefix:
 *
 *   unix:/tmp/foo.sock
 *   tcp::8438
 *   :8438                 - tcp listener
 */
func getListener(address string, socketPerms int) (listener net.Listener, err error) {
	if address == "" {
		return nil, fmt.Errorf("No listen address specified")
	}
	if idx := strings.Index(address, ":"); idx >= 0 {
		scheme := address[0:idx]
		switch scheme {
		case "tcp", "tcp4":
			path := address[idx+1:]
			if strings.Index(path, ":") == -1 {
				path = ":" + path
			}
			listener, err = net.Listen(scheme, path)
		case "tcp6": // general form: [host]:port
			path := address[idx+1:]
			if strings.Index(path, "[") != 0 { // port only
				if strings.Index(path, ":") != 0 { // no leading :
					path = ":" + path
				}
			}
			listener, err = net.Listen(scheme, path)
		case "unix", "unixpacket":
			path := address[idx+1:]
			if listener, err = net.Listen(scheme, path); err == nil {
				os.Chmod(path, os.FileMode(socketPerms))
			}
		default: // assume TCP4 address
			listener, err = net.Listen("tcp", address)
		}
	} else { // no scheme, port only specified
		listener, err = net.Listen("tcp", ":"+address)
	}
	return listener, err
}

/*
 * Handler for /scan
 *
 * Virus checks file and sends response
 */
func scanHandler(w http.ResponseWriter, req *http.Request) {
	if ctx.ShuttingDown {
		return
	}
	ctx.ActivityChan <- 1
	defer func() { ctx.ActivityChan <- -1 }()

	if !ctx.ScanInterceptor.Handle(w, req, req.Body) {
		w.Write([]byte("No virus found"))
	}
}

/*
 * Handler for scan & forward
 *
 * Constructs a forwarder and calls it
 */
func scanForwardHandler(w http.ResponseWriter, req *http.Request) {
	if ctx.ShuttingDown {
		return
	}
	ctx.ActivityChan <- 1
	defer func() { ctx.ActivityChan <- -1 }()

	fw := forwarder.NewForwarder(ctx.ApplicationURL, ctx.Config.App.ContentMemoryThreshold, ctx.ScanInterceptor)
	fw.SetLogger(ctx.Logger, ctx.Config.App.Debug)
	fw.HandleRequest(w, req)
}

/*
 * Handler for /info
 *
 * Validates the Scanner connection
 * Emits the information as a JSON response
 */
func infoHandler(w http.ResponseWriter, req *http.Request) {
	if ctx.ShuttingDown {
		return
	}
	ctx.ActivityChan <- 1
	defer func() { ctx.ActivityChan <- -1 }()

	info := &Info{
		Address: ctx.Scanner.Address(),
	}
	if err := ctx.Scanner.Ping(); err != nil {
		info.PingResult = err.Error()
	} else {
		info.PingResult = "Connected to server OK"
		if response, err := ctx.Scanner.Version(); err != nil {
			info.Version = err.Error()
		} else {
			info.Version = response
		}
		/*
		 * Validate the Clamd response for a viral string
		 */
		reader := bytes.NewReader(EICAR)
		if result, err := ctx.Scanner.Scan(reader); err != nil {
			info.TestScanVirusResult = err.Error()
		} else {
			info.TestScanVirusResult = result.String()
		}
		/*
		 * Validate the Clamd response for a non-viral string
		 */
		reader = bytes.NewReader([]byte("foo bar mcgrew"))
		if result, err := ctx.Scanner.Scan(reader); err != nil {
			info.TestScanCleanResult = err.Error()
		} else {
			info.TestScanCleanResult = result.String()
		}
	}
	// Aaaand return
	s, _ := json.Marshal(info)
	w.Header().Set("Content-Type", "application/json")
	w.Write([]byte(s))
}
<commit_msg>Remove DebugClam and associated side effects<commit_after>/*
 * The Clammit application intercepts HTTP POST requests with content-type
 * "multipart/form-data", forwards any "file" form-data elements to ClamAV
 * and only forwards the request to the application if ClamAV passes all
 * of these elements as virus-free.
 */
package main

import (
	"bytes"
	"clammit/forwarder"
	"clammit/scanner"
	"encoding/json"
	"flag"
	"fmt"
	"gopkg.in/gcfg.v1"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"runtime"
	"strconv"
	"strings"
	"syscall"
)

//
// Configuration structure, designed for gcfg
//
type Config struct {
	App ApplicationConfig `gcfg:"application"`
}

type ApplicationConfig struct {
	// The address to listen on. This can be one of:
	// * tcp:host:port
	// * tcp:port
	// * unix:filename
	// * host:port
	// * :port
	//
	// For example:
	//   Listen: tcp:0.0.0.0:8438
	//   Listen: unix:/tmp/clammit.sock
	//   Listen: :8438
	Listen string `gcfg:"listen"`
	// Socket file permissions (only used if listening on a unix socket), in octal form.
	//
	// For example:
	//   SocketPerms: 0766
	SocketPerms string `gcfg:"unix-socket-perms"`
	// The URL of the application that Clammit is proxying. Generally, this will
	// be the base URL (http://host:port/), but you can also add a path prefix
	// if needed (http://host:port/prefix)
	ApplicationURL string `gcfg:"application-url"`
	// The URL of clamd, which will either be TCP or Unix:
	//
	// For example:
	//   ClamdURL: tcp://localhost:3310
	//   ClamdURL: unix:/tmp/clamd.sock
	ClamdURL string `gcfg:"clamd-url"`
	// The HTTP status code to return when a virus is found
	VirusStatusCode int `gcfg:"virus-status-code"`
	// If the body content-length exceeds this value, it will be written to
	// disk. Below it, we'll hold the whole body in memory to improve speed.
	ContentMemoryThreshold int64 `gcfg:"content-memory-threshold"`
	// Log file name (default is to log to stdout)
	Logfile string `gcfg:"log-file"`
	// If true, clammit will expose a small test HTML page.
	TestPages bool `gcfg:"test-pages"`
	// If true, will log the progression of each request through the forwarder
	Debug bool `gcfg:"debug"`
	// Number of CPU threads to use
	NumThreads int `gcfg:"num-threads"`
}

//
// Default configuration
//
var DefaultApplicationConfig = ApplicationConfig{
	Listen:                 ":8438",
	SocketPerms:            "0777",
	ApplicationURL:         "",
	ClamdURL:               "",
	VirusStatusCode:        418,
	ContentMemoryThreshold: 1024 * 1024,
	Logfile:                "",
	TestPages:              true,
	Debug:                  false,
	NumThreads:             runtime.NumCPU(),
}

//
// Application context
//
type Ctx struct {
	Config          Config
	ApplicationURL  *url.URL
	ScanInterceptor *ScanInterceptor
	Scanner         scanner.Scanner
	Logger          *log.Logger
	Listener        net.Listener
	ActivityChan    chan int
	ShuttingDown    bool
}

//
// JSON server information response
//
type Info struct {
	Address             string `json:"scan_server_url"`
	PingResult          string `json:"ping_result"`
	Version             string `json:"version"`
	TestScanVirusResult string `json:"test_scan_virus"`
	TestScanCleanResult string `json:"test_scan_clean"`
}

//
// Global variables and config
//
var ctx *Ctx
var configFile string
var EICAR = []byte(`X5O!P%@AP[4\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*`)

func init() {
	flag.StringVar(&configFile, "config", "", "Configuration file")
}

func main() {
	/*
	 * Construct configuration, set up logging
	 */
	flag.Parse()
	ctx = &Ctx{
		ActivityChan: make(chan int),
		ShuttingDown: false,
	}

	ctx.Config.App = DefaultApplicationConfig

	if err := gcfg.ReadFileInto(&ctx.Config, configFile); err != nil {
		log.Fatalf("Configuration read failure: %s", err.Error())
	}

	// Socket perms are octal!
	socketPerms := 0777
	if ctx.Config.App.SocketPerms != "" {
		if sp, err := strconv.ParseInt(ctx.Config.App.SocketPerms, 8, 0); err == nil {
			socketPerms = int(sp)
		} else {
			log.Fatalf("SocketPerms invalid (expected 4-digit octal: %s", err.Error)
		}
	}

	// Allow multi-proc
	runtime.GOMAXPROCS(ctx.Config.App.NumThreads)

	startLogging()

	/*
	 * Construct objects, validate the URLs
	 */
	ctx.ApplicationURL = checkURL(ctx.Config.App.ApplicationURL)
	checkURL(ctx.Config.App.ClamdURL)

	ctx.Scanner = new(scanner.Clamav)
	ctx.Scanner.SetLogger(ctx.Logger, ctx.Config.App.Debug)
	ctx.Scanner.SetAddress(ctx.Config.App.ClamdURL)

	ctx.ScanInterceptor = &ScanInterceptor{
		VirusStatusCode: ctx.Config.App.VirusStatusCode,
		Scanner:         ctx.Scanner,
	}

	/*
	 * Set up the HTTP server
	 */
	router := http.NewServeMux()

	router.HandleFunc("/clammit", infoHandler)
	router.HandleFunc("/clammit/scan", scanHandler)
	if ctx.Config.App.TestPages {
		fs := http.FileServer(http.Dir("testfiles"))
		router.Handle("/clammit/test/", http.StripPrefix("/clammit/test/", fs))
	}
	router.HandleFunc("/", scanForwardHandler)

	if listener, err := getListener(ctx.Config.App.Listen, socketPerms); err != nil {
		ctx.Logger.Fatal("Unable to listen on: ", ctx.Config.App.Listen, ", reason: ", err)
	} else {
		ctx.Listener = listener
		beGraceful() // graceful shutdown from here on in
		ctx.Logger.Println("Listening on", ctx.Config.App.Listen)
		http.Serve(listener, router)
	}
}

/*
 * Starts logging
 */
func startLogging() {
	if ctx.Config.App.Logfile != "" {
		w, err := os.OpenFile(ctx.Config.App.Logfile, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0660)
		if err == nil {
			ctx.Logger = log.New(w, "", log.LstdFlags)
		} else {
			log.Fatal("Failed to open log file", ctx.Config.App.Logfile, ":", err)
		}
	} else {
		ctx.Logger = log.New(os.Stdout, "", log.LstdFlags)
		ctx.Logger.Println("No log file configured - using stdout")
	}
}

/*
 * Handles graceful shutdown. Sets ctx.ShuttingDown = true to stop any new
 * requests, then waits for active requests to complete before closing the
 * HTTP listener.
 */
func beGraceful() {
	sigchan := make(chan os.Signal)
	signal.Notify(sigchan, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		activity := 0
		for {
			select {
			case _ = <-sigchan:
				ctx.Logger.Println("Received termination signal")
				ctx.ShuttingDown = true
				for activity > 0 {
					ctx.Logger.Printf("There are %d active requests, waiting", activity)
					i := <-ctx.ActivityChan
					activity += i
				}
				// This will cause main() to continue from http.Serve()
				// it will also clean up the unix socket (if relevant)
				ctx.Listener.Close()
			case i := <-ctx.ActivityChan:
				activity += i
			}
		}
	}()
}

/*
 * Validates the URL is OK (fatal error if not) and returns it
 */
func checkURL(urlString string) *url.URL {
	parsedURL, err := url.Parse(urlString)
	if err != nil {
		log.Fatal("Invalid URL:", urlString)
	}
	return parsedURL
}

/*
 * Returns a TCP or Unix socket listener, according to the scheme prefix:
 *
 *   unix:/tmp/foo.sock
 *   tcp::8438
 *   :8438                 - tcp listener
 */
func getListener(address string, socketPerms int) (listener net.Listener, err error) {
	if address == "" {
		return nil, fmt.Errorf("No listen address specified")
	}
	if idx := strings.Index(address, ":"); idx >= 0 {
		scheme := address[0:idx]
		switch scheme {
		case "tcp", "tcp4":
			path := address[idx+1:]
			if strings.Index(path, ":") == -1 {
				path = ":" + path
			}
			listener, err = net.Listen(scheme, path)
		case "tcp6": // general form: [host]:port
			path := address[idx+1:]
			if strings.Index(path, "[") != 0 { // port only
				if strings.Index(path, ":") != 0 { // no leading :
					path = ":" + path
				}
			}
			listener, err = net.Listen(scheme, path)
		case "unix", "unixpacket":
			path := address[idx+1:]
			if listener, err = net.Listen(scheme, path); err == nil {
				os.Chmod(path, os.FileMode(socketPerms))
			}
		default: // assume TCP4 address
			listener, err = net.Listen("tcp", address)
		}
	} else { // no scheme, port only specified
		listener, err = net.Listen("tcp", ":"+address)
	}
	return listener, err
}

/*
 * Handler for /scan
 *
 * Virus checks file and sends response
 */
func scanHandler(w http.ResponseWriter, req *http.Request) {
	if ctx.ShuttingDown {
		return
	}
	ctx.ActivityChan <- 1
	defer func() { ctx.ActivityChan <- -1 }()

	if !ctx.ScanInterceptor.Handle(w, req, req.Body) {
		w.Write([]byte("No virus found"))
	}
}

/*
 * Handler for scan & forward
 *
 * Constructs a forwarder and calls it
 */
func scanForwardHandler(w http.ResponseWriter, req *http.Request) {
	if ctx.ShuttingDown {
		return
	}
	ctx.ActivityChan <- 1
	defer func() { ctx.ActivityChan <- -1 }()

	fw := forwarder.NewForwarder(ctx.ApplicationURL, ctx.Config.App.ContentMemoryThreshold, ctx.ScanInterceptor)
	fw.SetLogger(ctx.Logger, ctx.Config.App.Debug)
	fw.HandleRequest(w, req)
}

/*
 * Handler for /info
 *
 * Validates the Scanner connection
 * Emits the information as a JSON response
 */
func infoHandler(w http.ResponseWriter, req *http.Request) {
	if ctx.ShuttingDown {
		return
	}
	ctx.ActivityChan <- 1
	defer func() { ctx.ActivityChan <- -1 }()

	info := &Info{
		Address: ctx.Scanner.Address(),
	}
	if err := ctx.Scanner.Ping(); err != nil {
		info.PingResult = err.Error()
	} else {
		info.PingResult = "Connected to server OK"
		if response, err := ctx.Scanner.Version(); err != nil {
			info.Version = err.Error()
		} else {
			info.Version = response
		}
		/*
		 * Validate the Clamd response for a viral string
		 */
		reader := bytes.NewReader(EICAR)
		if result, err := ctx.Scanner.Scan(reader); err != nil {
			info.TestScanVirusResult = err.Error()
		} else {
			info.TestScanVirusResult = result.String()
		}
		/*
		 * Validate the Clamd response for a non-viral string
		 */
		reader = bytes.NewReader([]byte("foo bar mcgrew"))
		if result, err := ctx.Scanner.Scan(reader); err != nil {
			info.TestScanCleanResult = err.Error()
		} else {
			info.TestScanCleanResult = result.String()
		}
	}
	// Aaaand return
	s, _ := json.Marshal(info)
	w.Header().Set("Content-Type", "application/json")
	w.Write([]byte(s))
}
<|endoftext|>


<commit_before>/*
Package hamt32 implements a 32 node wide Hashed Array Mapped Trie. The hash key
is 30 bits wide and broken into 6 numbers of 5 bits each. Those 5bit numbers
allows us to index into a 32 node array. Each node is either a leaf or another
32 node table. So the 30bit hash allows us to index into a B+ Tree with a
branching factor of 32 and a Maximum depth of 6.

The basic insertion operation is to calculate a 30 bit hash value from your key
(a string in the case you use hamt.StringKey), then split it into six 5bit
 numbers. These six numbers represent a path thru the tree. For each level we
use the coresponding number as an index into the 32 cell array. If the cell is
empty we create a  leaf node there. If the cell is occupide by another table
we continue walking up the tree. If the cell is occupide by a leaf we promote
that cell to a new table and put the current leaf and new one into that table
in cells corresponding to that new level. If we are at the maximun depth of
the  tree and there is already a leaf there we insert our key,value pair into
that leaf.

The retrieval operation is a similar tree walk guided by the six 5bit numbers
till we find a leaf with the key,value pair in it.

The deletion operation is a walk to find the key, then delete the key from the
leaf. An empty leaf is removed from it's table. If the table has only one other
leaf in that level we will remove that leaf, replace the table in it's parent
table placing that last leaf down one level.

*/
package hamt32

import (
	"fmt"
	"strings"

	"github.com/lleo/go-hamt/key"
)

// nBits constant is the number of bits(5) a 30bit hash value is split into
// to provied the index of a HAMT.
const nBits uint = 5

// maxDepth constant is the maximum depth(5) of nBits values that constitute
// the path in a HAMT, from [0..maxDepth] for a total of maxDepth+1(6) levels.
// nBits*(maxDepth+1) == HASHBITS (ie 5*(5+1) == 30).
const maxDepth uint = 5

// tableCapacity constant is the number of table entries in a each node of
// a HAMT datastructure; its value is 2^5 == 32.
const tableCapacity uint = uint(1 << nBits)

// downgradeThreshold constant is the number of nodes a fullTable has shrunk to,
// before it is converted to a compressedTable.
const downgradeThreshold uint = tableCapacity / 8

// upgradeThreshold constan is the number of nodes a compressedTable has grown to,
// before it is converted to a fullTable.
const upgradeThreshold uint = tableCapacity / 2

func indexMask(depth uint) uint32 {
	return uint32(uint8(1<<nBits)-1) << (depth * nBits)
}

func index(h30 uint32, depth uint) uint {
	var idxMask = indexMask(depth)
	var idx = uint((h30 & idxMask) >> (depth * nBits))
	return idx
}

func hashPathString(hashPath uint32, depth uint) string {
	if depth == 0 {
		return "/"
	}
	var strs = make([]string, depth)

	for d := uint(0); d < depth; d++ {
		var idx = index(hashPath, d)
		strs[d] = fmt.Sprintf("%02d", idx)
	}

	return "/" + strings.Join(strs, "/")
}

func hash30String(h30 uint32) string {
	return hashPathString(h30, maxDepth)
}

func hashPathMask(depth uint) uint32 {
	return uint32(1<<(depth*nBits)) - 1
}

func buildHashPath(hashPath uint32, idx, depth uint) uint32 {
	return hashPath | uint32(idx<<(depth*nBits))
}

// Configuration contants to be passed to `hamt32.New(int) *Hamt`.
const (
	// HybridTables indicates the structure should use compressedTable
	// initially, then upgrad to fullTable when appropriate.
	HybridTables = iota
	// CompTablesOnly indicates the structure should use compressedTables ONLY.
	// This was intended just save space, but also seems to be faster; CPU cache
	// locality maybe?
	CompTablesOnly
	// FullTableOnly indicates the structure should use fullTables ONLY.
	// This was intended to be for speed, as compressed tables use a software
	// bitCount function to access individual cells. Turns out, not so much.
	FullTablesOnly
)

// TableOptionName is a map of the table option value Hybrid, CompTablesOnly,
// or FullTableOnly to a string representing that option.
//      var options = hamt32.FullTablesOnly
//      hamt32.TableOptionName[hamt32.FullTablesOnly] == "FullTablesOnly"
var TableOptionName = make(map[int]string, 3)

func init() {
	TableOptionName[HybridTables] = "HybridTables"
	TableOptionName[CompTablesOnly] = "CompTablesOnly"
	TableOptionName[FullTablesOnly] = "FullTablesOnly"
}

// Hamt is a Hashed Array Map Trie data structure. It has a branching factor of
// 32 and is at most 6 nodes deep. See:
// https://en.wikipedia.org/wiki/Hash_array_mapped_trie
type Hamt struct {
	root            tableI
	nentries        int
	grade, fullinit bool
}

// New creates a new hamt32.Hamt data structure with the table option set to
// either:
//
// `hamt32.HybridTables`:
// Initially start out with compressedTable, but when the table is half full
// upgrade to fullTable. If a fullTable shrinks to tableCapacity/8(4) entries
// downgrade to compressedTable.
//
// `hamt32.CompTablesOnly`:
// Use compressedTable ONLY with no up/downgrading to/from fullTable. This
// uses the least amount of space.
//
// `hamt32.FullTablesOnly`:
// Only use fullTable no up/downgrading from/to compressedTables. This is
// the fastest performance.
//
func New(opt int) *Hamt {
	var h = new(Hamt)
	if opt == CompTablesOnly {
		h.grade = false
		h.fullinit = false
	} else if opt == FullTablesOnly {
		h.grade = false
		h.fullinit = true
	} else /* opt == HybridTables */ {
		h.grade = true
		h.fullinit = false
	}
	return h
}

// IsEmpty Hamt method returns a boolean indicating if this Hamt structure has no entries.
func (h *Hamt) IsEmpty() bool {
	return h.root == nil
}

// Get Hamt method looks up a given key in the Hamt data structure.
func (h *Hamt) Get(k key.Key) (interface{}, bool) {
	if h.IsEmpty() {
		return nil, false
	}

	var h30 = k.Hash30()

	var curTable = h.root //ISA tableI

	for depth := uint(0); depth <= maxDepth; depth++ {
		var idx = index(h30, depth)
		var curNode = curTable.get(idx)

		if curNode == nil {
			break
		}

		if leaf, isLeaf := curNode.(leafI); isLeaf {
			var val, found = leaf.get(k)
			return val, found
		}

		//else curNode MUST BE A tableI
		curTable = curNode.(tableI)
	}
	// curNode == nil || depth > maxDepth

	return nil, false
}

// Put Hamt method inserts a given key/val pair into the Hamt data structure.
// It returns a boolean indicating if the key/val was inserted or whether or
// not the key already existed and the val was merely overwritten.
func (h *Hamt) Put(k key.Key, v interface{}) bool {
	var depth uint
	var hashPath uint32

	if h.IsEmpty() {
		h.root = h.newRootTable(depth, hashPath, newFlatLeaf(k, v))
		h.nentries++
		return true
	}

	var path = newPathT()
	var curTable = h.root

	for depth = 0; depth <= maxDepth; depth++ {
		var idx = index(k.Hash30(), depth)
		var curNode = curTable.get(idx)

		if curNode == nil {
			curTable.set(idx, newFlatLeaf(k, v))
			h.nentries++

			// upgrade?
			if h.grade {
				_, isCompressedTable := curTable.(*compressedTable)
				if isCompressedTable && curTable.nentries() >= upgradeThreshold {
					curTable = upgradeToFullTable(hashPath, curTable.entries())
					if depth == 0 {
						h.root = curTable
					} else {
						parentTable := path.peek()
						parentIdx := index(k.Hash30(), depth-1)
						parentTable.set(parentIdx, curTable)
					}
				}
			}

			return true //inserted
		}

		if curLeaf, isLeaf := curNode.(leafI); isLeaf {
			if curLeaf.Hash30() == k.Hash30() {
				// This is a minor optimization but since these two leaves
				// will collide all the way up the to maxDepth, we can
				// choose to create the collisionLeaf hear and now.

				// Accumulate collisionLeaf
				colLeaf, inserted := curLeaf.put(k, v)
				if inserted {
					curTable.set(idx, colLeaf)
					h.nentries++
				}
				return inserted
			}

			if depth == maxDepth {
				// this test should be delete cuz it is logically impossible
				if curLeaf.Hash30() != k.Hash30() {
					// This should not happen cuz we had to walk up maxDepth
					// levels to get here.
					panic("WTF!!!")
				}

				// Accumulate collisionLeaf
				insLeaf, inserted := curLeaf.put(k, v)
				if inserted {
					curTable.set(idx, insLeaf)
					h.nentries++
				}
				return inserted
			}

			hashPath = buildHashPath(hashPath, idx, depth)
			var collisionTable = h.newTable(depth+1, hashPath, curLeaf, newFlatLeaf(k, v))
			curTable.set(idx, collisionTable)
			h.nentries++

			return true
		}

		hashPath = buildHashPath(hashPath, idx, depth)
		path.push(curTable)
		curTable = curNode.(tableI)
	}

	//log.Println(path)
	//log.Printf("k=%s, v=%v", k, v)

	panic("WTF!!")
}

// Del Hamt Method removes a given key from the Hamt data structure. It returns
// a pair of values: the value stored and a boolean indicating if the key was
// even found and deleted.
func (h *Hamt) Del(k key.Key) (interface{}, bool) {
	if h.IsEmpty() {
		return nil, false
	}

	var h30 = k.Hash30()
	var depth uint
	var hashPath uint32

	var path = newPathT()
	var curTable = h.root

	for depth = 0; depth <= maxDepth; depth++ {
		var idx = index(h30, depth)
		var curNode = curTable.get(idx)

		if curNode == nil {
			return nil, false
		}

		if curLeaf, isLeaf := curNode.(leafI); isLeaf {
			val, delLeaf, deleted := curLeaf.del(k)
			if !deleted {
				return nil, false
			}
			// else a leaf key/value was deleted
			h.nentries--

			// If curLeaf was a collisionLeaf,
			// then delLeaf is either a slimmed down collisionLeaf or a flatLeaf.
			// If curLeaf was a flatLeaf then delLeaf is nil.
			curTable.set(idx, delLeaf)

			// downgrade?
			if h.grade {
				if delLeaf == nil {
					_, isFullTable := curTable.(*fullTable)
					if isFullTable && curTable.nentries() <= downgradeThreshold {
						curTable = downgradeToCompressedTable(hashPath, curTable.entries())
						if depth == 0 {
							h.root = curTable
						} else {
							parentTable := path.peek()
							parentIdx := index(h30, depth-1)
							parentTable.set(parentIdx, curTable)
						}
					}
				}
			}

			// If curTable has only one entry and that entry is a leafI,
			// then collapse that leafI down to the position curTable holds
			// in the parent Table; repeat test and collapse for parent table.
			//
			// These are identical for conditionals:
			//  curTable != h.root AND len(path) > 0 AND depth > 0
			//
			for curTable.nentries() == 1 && depth > 0 {
				// _ = ASSERT && Assert(curTable != h.root, "curTable == h.root")
				// _ = ASSERT && Assert(depth == len(path), "depth != len(path)")

				var node = curTable.entries()[0].node
				var leaf, isLeaf = node.(leafI)
				if !isLeaf {
					// We only collapse leafs
					break
				}

				// Collapse leaf down to where curTable is in parentTable

				var parentTable = path.pop()
				depth-- // OR depth = len(path)

				parentIdx := index(curTable.Hash30(), depth)
				parentTable.set(parentIdx, leaf)

				curTable = parentTable
			}

			// TODO: I should keep this table rather than throwing it away.
			// Instead using h.root == nil to detect emptyness, we should
			// trust our accounting and use h.nentries == 0.
			if curTable == h.root && curTable.nentries() == 0 {
				h.root = nil
			}

			return val, true
		} //if isLeaf

		// curNode is not nil
		// curNode is not a leafI
		// curNode MUST be a tableI
		hashPath = buildHashPath(hashPath, idx, depth)
		path.push(curTable)
		curTable = curNode.(tableI)
	} //for depth loop

	//log.Printf("WTF! this should never be called; k=%s", k)
	return nil, false
}

// String returns a string representation of the Hamt string.
func (h *Hamt) String() string {
	return fmt.Sprintf("Hamt{ nentries: %d, root: %s }", h.nentries, h.root.LongString("", 0))
}

// LongString returns a complete listing of the entire Hamt data structure.
func (h *Hamt) LongString(indent string) string {
	var str string
	if h.root != nil {
		str = indent + fmt.Sprintf("Hamt{ nentries: %d, root:\n", h.nentries)
		str += indent + h.root.LongString(indent, 0)
		str += indent + "} //Hamt"
	} else {
		str = indent + fmt.Sprintf("Hamt{ nentries: %d, root: nil }", h.nentries)
	}
	return str
}
<commit_msg>trivial comment indentation<commit_after>/*
Package hamt32 implements a 32 node wide Hashed Array Mapped Trie. The hash key
is 30 bits wide and broken into 6 numbers of 5 bits each. Those 5bit numbers
allows us to index into a 32 node array. Each node is either a leaf or another
32 node table. So the 30bit hash allows us to index into a B+ Tree with a
branching factor of 32 and a Maximum depth of 6.

The basic insertion operation is to calculate a 30 bit hash value from your key
(a string in the case you use hamt.StringKey), then split it into six 5bit
 numbers. These six numbers represent a path thru the tree. For each level we
use the coresponding number as an index into the 32 cell array. If the cell is
empty we create a  leaf node there. If the cell is occupide by another table
we continue walking up the tree. If the cell is occupide by a leaf we promote
that cell to a new table and put the current leaf and new one into that table
in cells corresponding to that new level. If we are at the maximun depth of
the  tree and there is already a leaf there we insert our key,value pair into
that leaf.

The retrieval operation is a similar tree walk guided by the six 5bit numbers
till we find a leaf with the key,value pair in it.

The deletion operation is a walk to find the key, then delete the key from the
leaf. An empty leaf is removed from it's table. If the table has only one other
leaf in that level we will remove that leaf, replace the table in it's parent
table placing that last leaf down one level.

*/
package hamt32

import (
	"fmt"
	"strings"

	"github.com/lleo/go-hamt/key"
)

// nBits constant is the number of bits(5) a 30bit hash value is split into
// to provied the index of a HAMT.
const nBits uint = 5

// maxDepth constant is the maximum depth(5) of nBits values that constitute
// the path in a HAMT, from [0..maxDepth] for a total of maxDepth+1(6) levels.
// nBits*(maxDepth+1) == HASHBITS (ie 5*(5+1) == 30).
const maxDepth uint = 5

// tableCapacity constant is the number of table entries in a each node of
// a HAMT datastructure; its value is 2^5 == 32.
const tableCapacity uint = uint(1 << nBits)

// downgradeThreshold constant is the number of nodes a fullTable has shrunk to,
// before it is converted to a compressedTable.
const downgradeThreshold uint = tableCapacity / 8

// upgradeThreshold constan is the number of nodes a compressedTable has grown to,
// before it is converted to a fullTable.
const upgradeThreshold uint = tableCapacity / 2

func indexMask(depth uint) uint32 {
	return uint32(uint8(1<<nBits)-1) << (depth * nBits)
}

func index(h30 uint32, depth uint) uint {
	var idxMask = indexMask(depth)
	var idx = uint((h30 & idxMask) >> (depth * nBits))
	return idx
}

func hashPathString(hashPath uint32, depth uint) string {
	if depth == 0 {
		return "/"
	}
	var strs = make([]string, depth)

	for d := uint(0); d < depth; d++ {
		var idx = index(hashPath, d)
		strs[d] = fmt.Sprintf("%02d", idx)
	}

	return "/" + strings.Join(strs, "/")
}

func hash30String(h30 uint32) string {
	return hashPathString(h30, maxDepth)
}

func hashPathMask(depth uint) uint32 {
	return uint32(1<<(depth*nBits)) - 1
}

func buildHashPath(hashPath uint32, idx, depth uint) uint32 {
	return hashPath | uint32(idx<<(depth*nBits))
}

// Configuration contants to be passed to `hamt32.New(int) *Hamt`.
const (
	// HybridTables indicates the structure should use compressedTable
	// initially, then upgrad to fullTable when appropriate.
	HybridTables = iota
	// CompTablesOnly indicates the structure should use compressedTables ONLY.
	// This was intended just save space, but also seems to be faster; CPU cache
	// locality maybe?
	CompTablesOnly
	// FullTableOnly indicates the structure should use fullTables ONLY.
	// This was intended to be for speed, as compressed tables use a software
	// bitCount function to access individual cells. Turns out, not so much.
	FullTablesOnly
)

// TableOptionName is a map of the table option value Hybrid, CompTablesOnly,
// or FullTableOnly to a string representing that option.
//      var options = hamt32.FullTablesOnly
//      hamt32.TableOptionName[hamt32.FullTablesOnly] == "FullTablesOnly"
var TableOptionName = make(map[int]string, 3)

func init() {
	TableOptionName[HybridTables] = "HybridTables"
	TableOptionName[CompTablesOnly] = "CompTablesOnly"
	TableOptionName[FullTablesOnly] = "FullTablesOnly"
}

//Hamt is a Hashed Array Map Trie data structure. It has a branching factor of
//32 and is at most 6 nodes deep. See:
//https://en.wikipedia.org/wiki/Hash_array_mapped_trie
type Hamt struct {
	root            tableI
	nentries        int
	grade, fullinit bool
}

//New creates a new hamt32.Hamt data structure with the table option set to
//either:
//
//`hamt32.HybridTables`:
//Initially start out with compressedTable, but when the table is half full
//upgrade to fullTable. If a fullTable shrinks to tableCapacity/8(4) entries
//downgrade to compressedTable.
//
//`hamt32.CompTablesOnly`:
//Use compressedTable ONLY with no up/downgrading to/from fullTable. This
//uses the least amount of space.
//
//`hamt32.FullTablesOnly`:
//Only use fullTable no up/downgrading from/to compressedTables. This is
//the fastest performance.
func New(opt int) *Hamt {
	var h = new(Hamt)
	if opt == CompTablesOnly {
		h.grade = false
		h.fullinit = false
	} else if opt == FullTablesOnly {
		h.grade = false
		h.fullinit = true
	} else /* opt == HybridTables */ {
		h.grade = true
		h.fullinit = false
	}
	return h
}

// IsEmpty Hamt method returns a boolean indicating if this Hamt structure has no entries.
func (h *Hamt) IsEmpty() bool {
	return h.root == nil
}

// Get Hamt method looks up a given key in the Hamt data structure.
func (h *Hamt) Get(k key.Key) (interface{}, bool) {
	if h.IsEmpty() {
		return nil, false
	}

	var h30 = k.Hash30()

	var curTable = h.root //ISA tableI

	for depth := uint(0); depth <= maxDepth; depth++ {
		var idx = index(h30, depth)
		var curNode = curTable.get(idx)

		if curNode == nil {
			break
		}

		if leaf, isLeaf := curNode.(leafI); isLeaf {
			var val, found = leaf.get(k)
			return val, found
		}

		//else curNode MUST BE A tableI
		curTable = curNode.(tableI)
	}
	// curNode == nil || depth > maxDepth

	return nil, false
}

// Put Hamt method inserts a given key/val pair into the Hamt data structure.
// It returns a boolean indicating if the key/val was inserted or whether or
// not the key already existed and the val was merely overwritten.
func (h *Hamt) Put(k key.Key, v interface{}) bool {
	var depth uint
	var hashPath uint32

	if h.IsEmpty() {
		h.root = h.newRootTable(depth, hashPath, newFlatLeaf(k, v))
		h.nentries++
		return true
	}

	var path = newPathT()
	var curTable = h.root

	for depth = 0; depth <= maxDepth; depth++ {
		var idx = index(k.Hash30(), depth)
		var curNode = curTable.get(idx)

		if curNode == nil {
			curTable.set(idx, newFlatLeaf(k, v))
			h.nentries++

			// upgrade?
			if h.grade {
				_, isCompressedTable := curTable.(*compressedTable)
				if isCompressedTable && curTable.nentries() >= upgradeThreshold {
					curTable = upgradeToFullTable(hashPath, curTable.entries())
					if depth == 0 {
						h.root = curTable
					} else {
						parentTable := path.peek()
						parentIdx := index(k.Hash30(), depth-1)
						parentTable.set(parentIdx, curTable)
					}
				}
			}

			return true //inserted
		}

		if curLeaf, isLeaf := curNode.(leafI); isLeaf {
			if curLeaf.Hash30() == k.Hash30() {
				// This is a minor optimization but since these two leaves
				// will collide all the way up the to maxDepth, we can
				// choose to create the collisionLeaf hear and now.

				// Accumulate collisionLeaf
				colLeaf, inserted := curLeaf.put(k, v)
				if inserted {
					curTable.set(idx, colLeaf)
					h.nentries++
				}
				return inserted
			}

			if depth == maxDepth {
				// this test should be delete cuz it is logically impossible
				if curLeaf.Hash30() != k.Hash30() {
					// This should not happen cuz we had to walk up maxDepth
					// levels to get here.
					panic("WTF!!!")
				}

				// Accumulate collisionLeaf
				insLeaf, inserted := curLeaf.put(k, v)
				if inserted {
					curTable.set(idx, insLeaf)
					h.nentries++
				}
				return inserted
			}

			hashPath = buildHashPath(hashPath, idx, depth)
			var collisionTable = h.newTable(depth+1, hashPath, curLeaf, newFlatLeaf(k, v))
			curTable.set(idx, collisionTable)
			h.nentries++

			return true
		}

		hashPath = buildHashPath(hashPath, idx, depth)
		path.push(curTable)
		curTable = curNode.(tableI)
	}

	//log.Println(path)
	//log.Printf("k=%s, v=%v", k, v)

	panic("WTF!!")
}

// Del Hamt Method removes a given key from the Hamt data structure. It returns
// a pair of values: the value stored and a boolean indicating if the key was
// even found and deleted.
func (h *Hamt) Del(k key.Key) (interface{}, bool) {
	if h.IsEmpty() {
		return nil, false
	}

	var h30 = k.Hash30()
	var depth uint
	var hashPath uint32

	var path = newPathT()
	var curTable = h.root

	for depth = 0; depth <= maxDepth; depth++ {
		var idx = index(h30, depth)
		var curNode = curTable.get(idx)

		if curNode == nil {
			return nil, false
		}

		if curLeaf, isLeaf := curNode.(leafI); isLeaf {
			val, delLeaf, deleted := curLeaf.del(k)
			if !deleted {
				return nil, false
			}
			// else a leaf key/value was deleted
			h.nentries--

			// If curLeaf was a collisionLeaf,
			// then delLeaf is either a slimmed down collisionLeaf or a flatLeaf.
			// If curLeaf was a flatLeaf then delLeaf is nil.
			curTable.set(idx, delLeaf)

			// downgrade?
			if h.grade {
				if delLeaf == nil {
					_, isFullTable := curTable.(*fullTable)
					if isFullTable && curTable.nentries() <= downgradeThreshold {
						curTable = downgradeToCompressedTable(hashPath, curTable.entries())
						if depth == 0 {
							h.root = curTable
						} else {
							parentTable := path.peek()
							parentIdx := index(h30, depth-1)
							parentTable.set(parentIdx, curTable)
						}
					}
				}
			}

			// If curTable has only one entry and that entry is a leafI,
			// then collapse that leafI down to the position curTable holds
			// in the parent Table; repeat test and collapse for parent table.
			//
			// These are identical for conditionals:
			//  curTable != h.root AND len(path) > 0 AND depth > 0
			//
			for curTable.nentries() == 1 && depth > 0 {
				// _ = ASSERT && Assert(curTable != h.root, "curTable == h.root")
				// _ = ASSERT && Assert(depth == len(path), "depth != len(path)")

				var node = curTable.entries()[0].node
				var leaf, isLeaf = node.(leafI)
				if !isLeaf {
					// We only collapse leafs
					break
				}

				// Collapse leaf down to where curTable is in parentTable

				var parentTable = path.pop()
				depth-- // OR depth = len(path)

				parentIdx := index(curTable.Hash30(), depth)
				parentTable.set(parentIdx, leaf)

				curTable = parentTable
			}

			// TODO: I should keep this table rather than throwing it away.
			// Instead using h.root == nil to detect emptyness, we should
			// trust our accounting and use h.nentries == 0.
			if curTable == h.root && curTable.nentries() == 0 {
				h.root = nil
			}

			return val, true
		} //if isLeaf

		// curNode is not nil
		// curNode is not a leafI
		// curNode MUST be a tableI
		hashPath = buildHashPath(hashPath, idx, depth)
		path.push(curTable)
		curTable = curNode.(tableI)
	} //for depth loop

	//log.Printf("WTF! this should never be called; k=%s", k)
	return nil, false
}

// String returns a string representation of the Hamt string.
func (h *Hamt) String() string {
	return fmt.Sprintf("Hamt{ nentries: %d, root: %s }", h.nentries, h.root.LongString("", 0))
}

// LongString returns a complete listing of the entire Hamt data structure.
func (h *Hamt) LongString(indent string) string {
	var str string
	if h.root != nil {
		str = indent + fmt.Sprintf("Hamt{ nentries: %d, root:\n", h.nentries)
		str += indent + h.root.LongString(indent, 0)
		str += indent + "} //Hamt"
	} else {
		str = indent + fmt.Sprintf("Hamt{ nentries: %d, root: nil }", h.nentries)
	}
	return str
}
<|endoftext|>