<commit_before>package bla

import (
	"encoding/json"
	"log"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"sort"
	"sync"
	"text/template"
	"time"

	"golang.org/x/net/webdav"

	"github.com/fsnotify/fsnotify"
)

// Initialize
type Handler struct {
	cfgPath string
	Cfg     *Config
	public  http.Handler
	webfs   http.Handler
	tpl     *template.Template

	publicPath   string
	templatePath string
	docPath      string

	mu       sync.RWMutex
	docs     map[string]*Doc
	sortDocs []*Doc
	tags     map[string][]*Doc
}

func NewHandler(cfgPath string) *Handler {

	h := &Handler{
		cfgPath: cfgPath,
		mu:      sync.RWMutex{},
	}

	h.loadConfig()
	h.watch()
	h.loadWebDav()

	return h
}

func (s *Handler) loadWebDav() {

	fs := webdav.Dir(s.Cfg.RootPath)
	ls := webdav.NewMemLS()

	handler := &webdav.Handler{
		Prefix:     "/fs",
		FileSystem: fs,
		LockSystem: ls,
	}
	s.webfs = handler
}

func (s *Handler) watch() {

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		log.Fatal(err)
	}

	go func() {
		// loadData minial interval is 1 second
		ticker := time.NewTicker(500 * time.Millisecond)
		mod := true

		for {
			select {
			case event := <-watcher.Events:
				switch ext := filepath.Ext(event.Name); ext {
				case ".md", ".json", ".tmpl":
					log.Println("modified file:", event.Name)
					mod = true
				}
			case err := <-watcher.Errors:
				log.Println("error:", err)
				return
			case <-ticker.C:
				if !mod {
					continue
				}
				mod = false
				s.loadData()
				s.loadTemplate()
				err := s.saveAll()
				if err != nil {
					log.Fatal("can't save docs:", err)
				}
			}
		}
	}()

	watcher.Add(s.docPath)
	watcher.Add(s.templatePath)
	watcher.Add(s.cfgPath)

	if err != nil {
		log.Print(err)
	}

}

func (s *Handler) saveAll() (err error) {

	log.Print("Saving all docs...")
	s.mu.Lock()
	defer s.mu.Unlock()

	log.Println("Deleting public...")
	err = os.RemoveAll(s.publicPath)
	if err != nil {
		return
	}

	var f *os.File
	err = os.MkdirAll(s.publicPath, 0755)
	if err != nil {
		return
	}

	log.Printf("saving all docs...")
	for _, doc := range s.docs {
		f, err = os.Create(filepath.Join(s.publicPath, doc.SlugTitle))
		if err != nil {
			return
		}
		if err = s.tpl.ExecuteTemplate(f, "single",
			&singleData{s, doc.Title, doc}); err != nil {
			return
		}
		f.Close()
	}

	var docs []*Doc
	if len(s.sortDocs) > s.Cfg.HomeDocCount {
		docs = s.sortDocs[:s.Cfg.HomeDocCount]
	} else {
		docs = s.sortDocs
	}

	f, err = os.Create(filepath.Join(s.publicPath, "index.html"))
	if err != nil {
		return
	}
	defer f.Close()

	if err = s.tpl.ExecuteTemplate(f, "index",
		&mulDocData{s, "", docs}); err != nil {
		return
	}

	log.Printf("saving all tags...")
	err = os.MkdirAll(filepath.Join(s.publicPath, "tags"), 0755)
	if err != nil {
		return
	}

	for tagName, docs := range s.tags {
		f, err = os.Create(filepath.Join(s.publicPath, "/tags/", tagName))
		if err != nil {
			return
		}
		sort.Sort(docsByTime(docs))
		if err = s.tpl.ExecuteTemplate(f, "tag_page",
			&tagData{s, tagName, docs, tagName}); err != nil {
			return
		}
		f.Close()
	}

	log.Printf("linking...%v", s.Cfg.LinkPath)
	for _, p := range s.Cfg.LinkPath {
		var realPath string
		realPath, err = filepath.Abs(p)
		if err != nil {
			return
		}
		err = os.Symlink(realPath, filepath.Join(s.publicPath, p))
		if err != nil {
			return
		}
	}
	log.Println("save completed")
	return nil
}

func (s *Handler) loadData() {
	log.Print("Loading docs from:", s.docPath)

	s.mu.Lock()
	s.sortDocs = []*Doc{}
	s.docs = map[string]*Doc{}
	s.tags = map[string][]*Doc{}
	s.mu.Unlock()

	f, err := os.Open(s.docPath)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	err = filepath.Walk(s.docPath, s.docWalker)
	if err != nil {
		log.Print(err)
	}
	sort.Sort(docsByTime(s.sortDocs))
	log.Print("End Loading docs from:", s.docPath)
}

func (s *Handler) docWalker(p string, info os.FileInfo, err error) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	start := time.Now()
	if info.IsDir() || filepath.Ext(info.Name()) != ".md" {
		return nil
	}
	var f *os.File
	f, err = os.Open(p)
	if err != nil {
		return err
	}
	defer f.Close()

	var doc *Doc
	doc, err = newDoc(f)
	if err != nil {
		return err
	}
	if !doc.Public {
		log.Printf("doc:%s loaded but not public", p)
		return nil
	}

	doc.SlugTitle = path.Base(p)[0 : len(path.Base(p))-3]

	for _, t := range doc.Tags {
		s.tags[t] = append(s.tags[t], doc)
	}
	s.docs[doc.SlugTitle] = doc
	s.sortDocs = append(s.sortDocs, doc)
	log.Printf("loaded doc:%s in %s", doc.SlugTitle,
		time.Now().Sub(start).String())
	return nil
}

func (h *Handler) loadConfig() {

	f, err := os.Open(h.cfgPath)
	if err != nil && os.IsExist(err) {
		log.Panic(err)
	}
	defer f.Close()

	log.Print("loading config")
	cfg := DefaultConfig()
	dec := json.NewDecoder(f)
	err = dec.Decode(cfg)

	if err != nil {
		log.Panic(err)
	}

	h.publicPath = filepath.Join(cfg.RootPath, "public")
	h.templatePath = filepath.Join(cfg.RootPath, "template")
	h.docPath = filepath.Join(cfg.RootPath, "docs")

	h.public = http.FileServer(http.Dir(h.publicPath))

	h.Cfg = cfg
	log.Printf("%#v", *cfg)

}

func (s *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {

	if r.URL.Path[:3] == "/fs" {
		s.webfs.ServeHTTP(w, r)
	} else {
		s.public.ServeHTTP(w, r)
	}
}

func (s *Handler) loadTemplate() {
	s.mu.Lock()
	defer s.mu.Unlock()

	log.Printf("loding template:%s", s.templatePath)
	tpl, err := template.ParseGlob(s.templatePath + "/*.tmpl")
	if err != nil {
		log.Print(err)
		return
	}
	s.tpl = tpl
}
<commit_msg>fix slice length<commit_after>package bla

import (
	"encoding/json"
	"log"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"text/template"
	"time"

	"golang.org/x/net/webdav"

	"github.com/fsnotify/fsnotify"
)

// Initialize
type Handler struct {
	cfgPath string
	Cfg     *Config
	public  http.Handler
	webfs   http.Handler
	tpl     *template.Template

	publicPath   string
	templatePath string
	docPath      string

	mu       sync.RWMutex
	docs     map[string]*Doc
	sortDocs []*Doc
	tags     map[string][]*Doc
}

func NewHandler(cfgPath string) *Handler {

	h := &Handler{
		cfgPath: cfgPath,
		mu:      sync.RWMutex{},
	}

	h.loadConfig()
	h.watch()
	h.loadWebDav()

	return h
}

func (s *Handler) loadWebDav() {

	fs := webdav.Dir(s.Cfg.RootPath)
	ls := webdav.NewMemLS()

	handler := &webdav.Handler{
		Prefix:     "/fs",
		FileSystem: fs,
		LockSystem: ls,
	}
	s.webfs = handler
}

func (s *Handler) watch() {

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		log.Fatal(err)
	}

	go func() {
		// loadData minial interval is 1 second
		ticker := time.NewTicker(500 * time.Millisecond)
		mod := true

		for {
			select {
			case event := <-watcher.Events:
				switch ext := filepath.Ext(event.Name); ext {
				case ".md", ".json", ".tmpl":
					log.Println("modified file:", event.Name)
					mod = true
				}
			case err := <-watcher.Errors:
				log.Println("error:", err)
				return
			case <-ticker.C:
				if !mod {
					continue
				}
				mod = false
				s.loadData()
				s.loadTemplate()
				err := s.saveAll()
				if err != nil {
					log.Fatal("can't save docs:", err)
				}
			}
		}
	}()

	watcher.Add(s.docPath)
	watcher.Add(s.templatePath)
	watcher.Add(s.cfgPath)

	if err != nil {
		log.Print(err)
	}

}

func (s *Handler) saveAll() (err error) {

	log.Print("Saving all docs...")
	s.mu.Lock()
	defer s.mu.Unlock()

	log.Println("Deleting public...")
	err = os.RemoveAll(s.publicPath)
	if err != nil {
		return
	}

	var f *os.File
	err = os.MkdirAll(s.publicPath, 0755)
	if err != nil {
		return
	}

	log.Printf("saving all docs...")
	for _, doc := range s.docs {
		f, err = os.Create(filepath.Join(s.publicPath, doc.SlugTitle))
		if err != nil {
			return
		}
		if err = s.tpl.ExecuteTemplate(f, "single",
			&singleData{s, doc.Title, doc}); err != nil {
			return
		}
		f.Close()
	}

	var docs []*Doc
	if len(s.sortDocs) > s.Cfg.HomeDocCount {
		docs = s.sortDocs[:s.Cfg.HomeDocCount]
	} else {
		docs = s.sortDocs
	}

	f, err = os.Create(filepath.Join(s.publicPath, "index.html"))
	if err != nil {
		return
	}
	defer f.Close()

	if err = s.tpl.ExecuteTemplate(f, "index",
		&mulDocData{s, "", docs}); err != nil {
		return
	}

	log.Printf("saving all tags...")
	err = os.MkdirAll(filepath.Join(s.publicPath, "tags"), 0755)
	if err != nil {
		return
	}

	for tagName, docs := range s.tags {
		f, err = os.Create(filepath.Join(s.publicPath, "/tags/", tagName))
		if err != nil {
			return
		}
		sort.Sort(docsByTime(docs))
		if err = s.tpl.ExecuteTemplate(f, "tag_page",
			&tagData{s, tagName, docs, tagName}); err != nil {
			return
		}
		f.Close()
	}

	log.Printf("linking...%v", s.Cfg.LinkPath)
	for _, p := range s.Cfg.LinkPath {
		var realPath string
		realPath, err = filepath.Abs(p)
		if err != nil {
			return
		}
		err = os.Symlink(realPath, filepath.Join(s.publicPath, p))
		if err != nil {
			return
		}
	}
	log.Println("save completed")
	return nil
}

func (s *Handler) loadData() {
	log.Print("Loading docs from:", s.docPath)

	s.mu.Lock()
	s.sortDocs = []*Doc{}
	s.docs = map[string]*Doc{}
	s.tags = map[string][]*Doc{}
	s.mu.Unlock()

	f, err := os.Open(s.docPath)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	err = filepath.Walk(s.docPath, s.docWalker)
	if err != nil {
		log.Print(err)
	}
	sort.Sort(docsByTime(s.sortDocs))
	log.Print("End Loading docs from:", s.docPath)
}

func (s *Handler) docWalker(p string, info os.FileInfo, err error) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	start := time.Now()
	if info.IsDir() || filepath.Ext(info.Name()) != ".md" {
		return nil
	}
	var f *os.File
	f, err = os.Open(p)
	if err != nil {
		return err
	}
	defer f.Close()

	var doc *Doc
	doc, err = newDoc(f)
	if err != nil {
		return err
	}
	if !doc.Public {
		log.Printf("doc:%s loaded but not public", p)
		return nil
	}

	doc.SlugTitle = path.Base(p)[0 : len(path.Base(p))-3]

	for _, t := range doc.Tags {
		s.tags[t] = append(s.tags[t], doc)
	}
	s.docs[doc.SlugTitle] = doc
	s.sortDocs = append(s.sortDocs, doc)
	log.Printf("loaded doc:%s in %s", doc.SlugTitle,
		time.Now().Sub(start).String())
	return nil
}

func (h *Handler) loadConfig() {

	f, err := os.Open(h.cfgPath)
	if err != nil && os.IsExist(err) {
		log.Panic(err)
	}
	defer f.Close()

	log.Print("loading config")
	cfg := DefaultConfig()
	dec := json.NewDecoder(f)
	err = dec.Decode(cfg)

	if err != nil {
		log.Panic(err)
	}

	h.publicPath = filepath.Join(cfg.RootPath, "public")
	h.templatePath = filepath.Join(cfg.RootPath, "template")
	h.docPath = filepath.Join(cfg.RootPath, "docs")

	h.public = http.FileServer(http.Dir(h.publicPath))

	h.Cfg = cfg
	log.Printf("%#v", *cfg)

}

func (s *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {

	if strings.HasPrefix(r.URL.Path, "/fs") {
		s.webfs.ServeHTTP(w, r)
	} else {
		s.public.ServeHTTP(w, r)
	}
}

func (s *Handler) loadTemplate() {
	s.mu.Lock()
	defer s.mu.Unlock()

	log.Printf("loding template:%s", s.templatePath)
	tpl, err := template.ParseGlob(s.templatePath + "/*.tmpl")
	if err != nil {
		log.Print(err)
		return
	}
	s.tpl = tpl
}
<|endoftext|>


<commit_before>// +build linux

package certs

import (
	"crypto/x509"
	"fmt"
	// "os/exec"
)

func FindCerts() ([]*x509.Certificate, error) {
	fmt.Println("find linux")

	return nil, nil
}
<commit_msg>certs/linux: Find certs on linux<commit_after>// +build linux

package certs

import (
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	// "os/exec"
)

// From Go's source, src/crypto/x509/root_linux.go
var certFiles = []string {
	"/etc/ssl/certs/ca-certificates.crt",                // Debian/Ubuntu/Gentoo etc.
	"/etc/pki/ca-trust/extracted/pem/tls-ca-bundle.pem", // CentOS/RHEL 7
	"/etc/pki/tls/certs/ca-bundle.crt",                  // Fedora/RHEL 6
	"/etc/ssl/ca-bundle.pem",                            // OpenSUSE
	"/etc/pki/tls/cacert.pem",                           // OpenELEC
}

func FindCerts() ([]*x509.Certificate, error) {
	for i := range certFiles {
		path, err := filepath.Abs(certFiles[i])

		fmt.Printf("checking %s\n", certFiles[i])

		if err != nil {
			return nil, err
		}
		_, err = os.Stat(path)
		if err == nil {
			bytes, err := ioutil.ReadFile(path)
			if err != nil {
				return nil, err
			}
			certs, err := ParsePEMIntoCerts(bytes)
			if err != nil {
				return nil, err
			}
			return certs, nil
		}
	}
	return nil, nil
}
<|endoftext|>


<commit_before>package lnd

import (
	"fmt"
	"net"

	"github.com/btcsuite/btcd/btcec"
	"github.com/btcsuite/btcd/wire"
	"github.com/lightningnetwork/lnd/chanbackup"
	"github.com/lightningnetwork/lnd/channelnotifier"
)

// addrSource is an interface that allow us to get the addresses for a target
// node. We'll need this in order to be able to properly proxy the
// notifications to create SCBs.
type addrSource interface {
	// AddrsForNode returns all known addresses for the target node public
	// key.
	AddrsForNode(nodePub *btcec.PublicKey) ([]net.Addr, error)
}

// channelNotifier is an implementation of the chanbackup.ChannelNotifier
// interface using the existing channelnotifier.ChannelNotifier struct. This
// implementation allows us to satisfy all the dependencies of the
// chanbackup.SubSwapper struct.
type channelNotifier struct {
	// chanNotifier is the based channel notifier that we'll proxy requests
	// from.
	chanNotifier *channelnotifier.ChannelNotifier

	// addrs is an implementation of the addrSource interface that allows
	// us to get the latest set of addresses for a given node. We'll need
	// this to be able to create an SCB for new channels.
	addrs addrSource
}

// SubscribeChans requests a new channel subscription relative to the initial
// set of known channels. We use the knownChans as a synchronization point to
// ensure that the chanbackup.SubSwapper does not miss any channel open or
// close events in the period between when it's created, and when it requests
// the channel subscription.
//
// NOTE: This is part of the chanbackup.ChannelNotifier interface.
func (c *channelNotifier) SubscribeChans(startingChans map[wire.OutPoint]struct{}) (
	*chanbackup.ChannelSubscription, error) {

	ltndLog.Infof("Channel backup proxy channel notifier starting")

	// TODO(roasbeef): read existing set of chans and diff

	quit := make(chan struct{})
	chanUpdates := make(chan chanbackup.ChannelEvent, 1)

	// In order to adhere to the interface, we'll proxy the events from the
	// channel notifier to the sub-swapper in a format it understands.
	go func() {
		// First, we'll subscribe to the primary channel notifier so we can
		// obtain events for new opened/closed channels.
		chanSubscription, err := c.chanNotifier.SubscribeChannelEvents()
		if err != nil {
			panic(fmt.Sprintf("unable to subscribe to chans: %v",
				err))
		}

		defer chanSubscription.Cancel()

		for {
			select {

			// A new event has been sent by the chanNotifier, we'll
			// filter out the events we actually care about and
			// send them to the sub-swapper.
			case e := <-chanSubscription.Updates():
				// TODO(roasbeef): batch dispatch ntnfs

				switch event := e.(type) {

				// A new channel has been opened, we'll obtain
				// the node address, then send to the
				// sub-swapper.
				case channelnotifier.OpenChannelEvent:
					nodeAddrs, err := c.addrs.AddrsForNode(
						event.Channel.IdentityPub,
					)
					if err != nil {
						pub := event.Channel.IdentityPub
						ltndLog.Errorf("unable to "+
							"fetch addrs for %x: %v",
							pub.SerializeCompressed(),
							err)
					}

					channel := event.Channel
					chanEvent := chanbackup.ChannelEvent{
						NewChans: []chanbackup.ChannelWithAddrs{
							{
								OpenChannel: channel,
								Addrs:       nodeAddrs,
							},
						},
					}

					select {
					case chanUpdates <- chanEvent:
					case <-quit:
						return
					}

				// An existing channel has been closed, we'll
				// send only the chanPoint of the closed
				// channel to the sub-swapper.
				case channelnotifier.ClosedChannelEvent:
					chanPoint := event.CloseSummary.ChanPoint
					chanEvent := chanbackup.ChannelEvent{
						ClosedChans: []wire.OutPoint{
							chanPoint,
						},
					}

					select {
					case chanUpdates <- chanEvent:
					case <-quit:
						return
					}
				}

			// The cancel method has been called, signalling us to
			// exit
			case <-quit:
				return
			}
		}
	}()

	return &chanbackup.ChannelSubscription{
		ChanUpdates: chanUpdates,
		Cancel: func() {
			close(quit)
		},
	}, nil
}

// A compile-time constraint to ensure channelNotifier implements
// chanbackup.ChannelNotifier.
var _ chanbackup.ChannelNotifier = (*channelNotifier)(nil)
<commit_msg>channel_notifier: update backup on pending chan<commit_after>package lnd

import (
	"fmt"
	"net"

	"github.com/btcsuite/btcd/btcec"
	"github.com/btcsuite/btcd/wire"
	"github.com/lightningnetwork/lnd/chanbackup"
	"github.com/lightningnetwork/lnd/channeldb"
	"github.com/lightningnetwork/lnd/channelnotifier"
)

// addrSource is an interface that allow us to get the addresses for a target
// node. We'll need this in order to be able to properly proxy the
// notifications to create SCBs.
type addrSource interface {
	// AddrsForNode returns all known addresses for the target node public
	// key.
	AddrsForNode(nodePub *btcec.PublicKey) ([]net.Addr, error)
}

// channelNotifier is an implementation of the chanbackup.ChannelNotifier
// interface using the existing channelnotifier.ChannelNotifier struct. This
// implementation allows us to satisfy all the dependencies of the
// chanbackup.SubSwapper struct.
type channelNotifier struct {
	// chanNotifier is the based channel notifier that we'll proxy requests
	// from.
	chanNotifier *channelnotifier.ChannelNotifier

	// addrs is an implementation of the addrSource interface that allows
	// us to get the latest set of addresses for a given node. We'll need
	// this to be able to create an SCB for new channels.
	addrs addrSource
}

// SubscribeChans requests a new channel subscription relative to the initial
// set of known channels. We use the knownChans as a synchronization point to
// ensure that the chanbackup.SubSwapper does not miss any channel open or
// close events in the period between when it's created, and when it requests
// the channel subscription.
//
// NOTE: This is part of the chanbackup.ChannelNotifier interface.
func (c *channelNotifier) SubscribeChans(startingChans map[wire.OutPoint]struct{}) (
	*chanbackup.ChannelSubscription, error) {

	ltndLog.Infof("Channel backup proxy channel notifier starting")

	// TODO(roasbeef): read existing set of chans and diff

	quit := make(chan struct{})
	chanUpdates := make(chan chanbackup.ChannelEvent, 1)

	// sendChanOpenUpdate is a closure that sends a ChannelEvent to the
	// chanUpdates channel to inform subscribers about new pending or
	// confirmed channels.
	sendChanOpenUpdate := func(newOrPendingChan *channeldb.OpenChannel) {
		nodeAddrs, err := c.addrs.AddrsForNode(
			newOrPendingChan.IdentityPub,
		)
		if err != nil {
			pub := newOrPendingChan.IdentityPub
			ltndLog.Errorf("unable to fetch addrs for %x: %v",
				pub.SerializeCompressed(), err)
		}

		chanEvent := chanbackup.ChannelEvent{
			NewChans: []chanbackup.ChannelWithAddrs{
				{
					OpenChannel: newOrPendingChan,
					Addrs:       nodeAddrs,
				},
			},
		}

		select {
		case chanUpdates <- chanEvent:
		case <-quit:
			return
		}
	}

	// In order to adhere to the interface, we'll proxy the events from the
	// channel notifier to the sub-swapper in a format it understands.
	go func() {
		// First, we'll subscribe to the primary channel notifier so we can
		// obtain events for new opened/closed channels.
		chanSubscription, err := c.chanNotifier.SubscribeChannelEvents()
		if err != nil {
			panic(fmt.Sprintf("unable to subscribe to chans: %v",
				err))
		}

		defer chanSubscription.Cancel()

		for {
			select {

			// A new event has been sent by the chanNotifier, we'll
			// filter out the events we actually care about and
			// send them to the sub-swapper.
			case e := <-chanSubscription.Updates():
				// TODO(roasbeef): batch dispatch ntnfs

				switch event := e.(type) {
				// A new channel has been opened and is still
				// pending. We can still create a backup, even
				// if the final channel ID is not yet available.
				case channelnotifier.PendingOpenChannelEvent:
					pendingChan := event.PendingChannel
					sendChanOpenUpdate(pendingChan)

				// A new channel has been confirmed, we'll
				// obtain the node address, then send to the
				// sub-swapper.
				case channelnotifier.OpenChannelEvent:
					sendChanOpenUpdate(event.Channel)

				// An existing channel has been closed, we'll
				// send only the chanPoint of the closed
				// channel to the sub-swapper.
				case channelnotifier.ClosedChannelEvent:
					chanPoint := event.CloseSummary.ChanPoint
					chanEvent := chanbackup.ChannelEvent{
						ClosedChans: []wire.OutPoint{
							chanPoint,
						},
					}

					select {
					case chanUpdates <- chanEvent:
					case <-quit:
						return
					}
				}

			// The cancel method has been called, signalling us to
			// exit
			case <-quit:
				return
			}
		}
	}()

	return &chanbackup.ChannelSubscription{
		ChanUpdates: chanUpdates,
		Cancel: func() {
			close(quit)
		},
	}, nil
}

// A compile-time constraint to ensure channelNotifier implements
// chanbackup.ChannelNotifier.
var _ chanbackup.ChannelNotifier = (*channelNotifier)(nil)
<|endoftext|>


<commit_before>package tarexport

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"

	"github.com/Sirupsen/logrus"
	"github.com/docker/distribution"
	"github.com/docker/docker/image"
	"github.com/docker/docker/image/v1"
	"github.com/docker/docker/layer"
	"github.com/docker/docker/pkg/archive"
	"github.com/docker/docker/pkg/chrootarchive"
	"github.com/docker/docker/pkg/progress"
	"github.com/docker/docker/pkg/streamformatter"
	"github.com/docker/docker/pkg/stringid"
	"github.com/docker/docker/pkg/symlink"
	"github.com/docker/docker/reference"
)

func (l *tarexporter) Load(inTar io.ReadCloser, outStream io.Writer, quiet bool) error {
	var (
		sf             = streamformatter.NewJSONStreamFormatter()
		progressOutput progress.Output
	)
	if !quiet {
		progressOutput = sf.NewProgressOutput(outStream, false)
		outStream = &streamformatter.StdoutFormatter{Writer: outStream, StreamFormatter: streamformatter.NewJSONStreamFormatter()}
	}

	tmpDir, err := ioutil.TempDir("", "docker-import-")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tmpDir)

	if err := chrootarchive.Untar(inTar, tmpDir, nil); err != nil {
		return err
	}
	// read manifest, if no file then load in legacy mode
	manifestPath, err := safePath(tmpDir, manifestFileName)
	if err != nil {
		return err
	}
	manifestFile, err := os.Open(manifestPath)
	if err != nil {
		if os.IsNotExist(err) {
			return l.legacyLoad(tmpDir, outStream, progressOutput)
		}
		return manifestFile.Close()
	}
	defer manifestFile.Close()

	var manifest []manifestItem
	if err := json.NewDecoder(manifestFile).Decode(&manifest); err != nil {
		return err
	}

	var parentLinks []parentLink
	var imageIDsStr string
	var imageRefCount int

	for _, m := range manifest {
		configPath, err := safePath(tmpDir, m.Config)
		if err != nil {
			return err
		}
		config, err := ioutil.ReadFile(configPath)
		if err != nil {
			return err
		}
		img, err := image.NewFromJSON(config)
		if err != nil {
			return err
		}
		var rootFS image.RootFS
		rootFS = *img.RootFS
		rootFS.DiffIDs = nil

		if expected, actual := len(m.Layers), len(img.RootFS.DiffIDs); expected != actual {
			return fmt.Errorf("invalid manifest, layers length mismatch: expected %q, got %q", expected, actual)
		}

		for i, diffID := range img.RootFS.DiffIDs {
			layerPath, err := safePath(tmpDir, m.Layers[i])
			if err != nil {
				return err
			}
			r := rootFS
			r.Append(diffID)
			newLayer, err := l.ls.Get(r.ChainID())
			if err != nil {
				newLayer, err = l.loadLayer(layerPath, rootFS, diffID.String(), m.LayerSources[diffID], progressOutput)
				if err != nil {
					return err
				}
			}
			defer layer.ReleaseAndLog(l.ls, newLayer)
			if expected, actual := diffID, newLayer.DiffID(); expected != actual {
				return fmt.Errorf("invalid diffID for layer %d: expected %q, got %q", i, expected, actual)
			}
			rootFS.Append(diffID)
		}

		imgID, err := l.is.Create(config)
		if err != nil {
			return err
		}
		imageIDsStr += fmt.Sprintf("Loaded image ID: %s\n", imgID)

		imageRefCount = 0
		for _, repoTag := range m.RepoTags {
			named, err := reference.ParseNamed(repoTag)
			if err != nil {
				return err
			}
			ref, ok := named.(reference.NamedTagged)
			if !ok {
				return fmt.Errorf("invalid tag %q", repoTag)
			}
			l.setLoadedTag(ref, imgID, outStream)
			outStream.Write([]byte(fmt.Sprintf("Loaded image: %s\n", ref)))
			imageRefCount++
		}

		parentLinks = append(parentLinks, parentLink{imgID, m.Parent})
		l.loggerImgEvent.LogImageEvent(imgID.String(), imgID.String(), "load")
	}

	for _, p := range validatedParentLinks(parentLinks) {
		if p.parentID != "" {
			if err := l.setParentID(p.id, p.parentID); err != nil {
				return err
			}
		}
	}

	if imageRefCount == 0 {
		outStream.Write([]byte(imageIDsStr))
	}

	return nil
}

func (l *tarexporter) setParentID(id, parentID image.ID) error {
	img, err := l.is.Get(id)
	if err != nil {
		return err
	}
	parent, err := l.is.Get(parentID)
	if err != nil {
		return err
	}
	if !checkValidParent(img, parent) {
		return fmt.Errorf("image %v is not a valid parent for %v", parent.ID(), img.ID())
	}
	return l.is.SetParent(id, parentID)
}

func (l *tarexporter) loadLayer(filename string, rootFS image.RootFS, id string, foreignSrc distribution.Descriptor, progressOutput progress.Output) (layer.Layer, error) {
	rawTar, err := os.Open(filename)
	if err != nil {
		logrus.Debugf("Error reading embedded tar: %v", err)
		return nil, err
	}
	defer rawTar.Close()

	inflatedLayerData, err := archive.DecompressStream(rawTar)
	if err != nil {
		return nil, err
	}
	defer inflatedLayerData.Close()

	if progressOutput != nil {
		fileInfo, err := os.Stat(filename)
		if err != nil {
			logrus.Debugf("Error statting file: %v", err)
			return nil, err
		}

		progressReader := progress.NewProgressReader(inflatedLayerData, progressOutput, fileInfo.Size(), stringid.TruncateID(id), "Loading layer")

		if ds, ok := l.ls.(layer.DescribableStore); ok {
			return ds.RegisterWithDescriptor(progressReader, rootFS.ChainID(), foreignSrc)
		}
		return l.ls.Register(progressReader, rootFS.ChainID())

	}

	if ds, ok := l.ls.(layer.DescribableStore); ok {
		return ds.RegisterWithDescriptor(inflatedLayerData, rootFS.ChainID(), foreignSrc)
	}
	return l.ls.Register(inflatedLayerData, rootFS.ChainID())
}

func (l *tarexporter) setLoadedTag(ref reference.NamedTagged, imgID image.ID, outStream io.Writer) error {
	if prevID, err := l.rs.Get(ref); err == nil && prevID != imgID {
		fmt.Fprintf(outStream, "The image %s already exists, renaming the old one with ID %s to empty string\n", ref.String(), string(prevID)) // todo: this message is wrong in case of multiple tags
	}

	if err := l.rs.AddTag(ref, imgID, true); err != nil {
		return err
	}
	return nil
}

func (l *tarexporter) legacyLoad(tmpDir string, outStream io.Writer, progressOutput progress.Output) error {
	legacyLoadedMap := make(map[string]image.ID)

	dirs, err := ioutil.ReadDir(tmpDir)
	if err != nil {
		return err
	}

	// every dir represents an image
	for _, d := range dirs {
		if d.IsDir() {
			if err := l.legacyLoadImage(d.Name(), tmpDir, legacyLoadedMap, progressOutput); err != nil {
				return err
			}
		}
	}

	// load tags from repositories file
	repositoriesPath, err := safePath(tmpDir, legacyRepositoriesFileName)
	if err != nil {
		return err
	}
	repositoriesFile, err := os.Open(repositoriesPath)
	if err != nil {
		if !os.IsNotExist(err) {
			return err
		}
		return repositoriesFile.Close()
	}
	defer repositoriesFile.Close()

	repositories := make(map[string]map[string]string)
	if err := json.NewDecoder(repositoriesFile).Decode(&repositories); err != nil {
		return err
	}

	for name, tagMap := range repositories {
		for tag, oldID := range tagMap {
			imgID, ok := legacyLoadedMap[oldID]
			if !ok {
				return fmt.Errorf("invalid target ID: %v", oldID)
			}
			named, err := reference.WithName(name)
			if err != nil {
				return err
			}
			ref, err := reference.WithTag(named, tag)
			if err != nil {
				return err
			}
			l.setLoadedTag(ref, imgID, outStream)
		}
	}

	return nil
}

func (l *tarexporter) legacyLoadImage(oldID, sourceDir string, loadedMap map[string]image.ID, progressOutput progress.Output) error {
	if _, loaded := loadedMap[oldID]; loaded {
		return nil
	}
	configPath, err := safePath(sourceDir, filepath.Join(oldID, legacyConfigFileName))
	if err != nil {
		return err
	}
	imageJSON, err := ioutil.ReadFile(configPath)
	if err != nil {
		logrus.Debugf("Error reading json: %v", err)
		return err
	}

	var img struct{ Parent string }
	if err := json.Unmarshal(imageJSON, &img); err != nil {
		return err
	}

	var parentID image.ID
	if img.Parent != "" {
		for {
			var loaded bool
			if parentID, loaded = loadedMap[img.Parent]; !loaded {
				if err := l.legacyLoadImage(img.Parent, sourceDir, loadedMap, progressOutput); err != nil {
					return err
				}
			} else {
				break
			}
		}
	}

	// todo: try to connect with migrate code
	rootFS := image.NewRootFS()
	var history []image.History

	if parentID != "" {
		parentImg, err := l.is.Get(parentID)
		if err != nil {
			return err
		}

		rootFS = parentImg.RootFS
		history = parentImg.History
	}

	layerPath, err := safePath(sourceDir, filepath.Join(oldID, legacyLayerFileName))
	if err != nil {
		return err
	}
	newLayer, err := l.loadLayer(layerPath, *rootFS, oldID, distribution.Descriptor{}, progressOutput)
	if err != nil {
		return err
	}
	rootFS.Append(newLayer.DiffID())

	h, err := v1.HistoryFromConfig(imageJSON, false)
	if err != nil {
		return err
	}
	history = append(history, h)

	config, err := v1.MakeConfigFromV1Config(imageJSON, rootFS, history)
	if err != nil {
		return err
	}
	imgID, err := l.is.Create(config)
	if err != nil {
		return err
	}

	metadata, err := l.ls.Release(newLayer)
	layer.LogReleaseMetadata(metadata)
	if err != nil {
		return err
	}

	if parentID != "" {
		if err := l.is.SetParent(imgID, parentID); err != nil {
			return err
		}
	}

	loadedMap[oldID] = imgID
	return nil
}

func safePath(base, path string) (string, error) {
	return symlink.FollowSymlinkInScope(filepath.Join(base, path), base)
}

type parentLink struct {
	id, parentID image.ID
}

func validatedParentLinks(pl []parentLink) (ret []parentLink) {
mainloop:
	for i, p := range pl {
		ret = append(ret, p)
		for _, p2 := range pl {
			if p2.id == p.parentID && p2.id != p.id {
				continue mainloop
			}
		}
		ret[i].parentID = ""
	}
	return
}

func checkValidParent(img, parent *image.Image) bool {
	if len(img.History) == 0 && len(parent.History) == 0 {
		return true // having history is not mandatory
	}
	if len(img.History)-len(parent.History) != 1 {
		return false
	}
	for i, h := range parent.History {
		if !reflect.DeepEqual(h, img.History[i]) {
			return false
		}
	}
	return true
}
<commit_msg>Report load progress on compressed layer size<commit_after>package tarexport

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"

	"github.com/Sirupsen/logrus"
	"github.com/docker/distribution"
	"github.com/docker/docker/image"
	"github.com/docker/docker/image/v1"
	"github.com/docker/docker/layer"
	"github.com/docker/docker/pkg/archive"
	"github.com/docker/docker/pkg/chrootarchive"
	"github.com/docker/docker/pkg/progress"
	"github.com/docker/docker/pkg/streamformatter"
	"github.com/docker/docker/pkg/stringid"
	"github.com/docker/docker/pkg/symlink"
	"github.com/docker/docker/reference"
)

func (l *tarexporter) Load(inTar io.ReadCloser, outStream io.Writer, quiet bool) error {
	var (
		sf             = streamformatter.NewJSONStreamFormatter()
		progressOutput progress.Output
	)
	if !quiet {
		progressOutput = sf.NewProgressOutput(outStream, false)
		outStream = &streamformatter.StdoutFormatter{Writer: outStream, StreamFormatter: streamformatter.NewJSONStreamFormatter()}
	}

	tmpDir, err := ioutil.TempDir("", "docker-import-")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tmpDir)

	if err := chrootarchive.Untar(inTar, tmpDir, nil); err != nil {
		return err
	}
	// read manifest, if no file then load in legacy mode
	manifestPath, err := safePath(tmpDir, manifestFileName)
	if err != nil {
		return err
	}
	manifestFile, err := os.Open(manifestPath)
	if err != nil {
		if os.IsNotExist(err) {
			return l.legacyLoad(tmpDir, outStream, progressOutput)
		}
		return manifestFile.Close()
	}
	defer manifestFile.Close()

	var manifest []manifestItem
	if err := json.NewDecoder(manifestFile).Decode(&manifest); err != nil {
		return err
	}

	var parentLinks []parentLink
	var imageIDsStr string
	var imageRefCount int

	for _, m := range manifest {
		configPath, err := safePath(tmpDir, m.Config)
		if err != nil {
			return err
		}
		config, err := ioutil.ReadFile(configPath)
		if err != nil {
			return err
		}
		img, err := image.NewFromJSON(config)
		if err != nil {
			return err
		}
		var rootFS image.RootFS
		rootFS = *img.RootFS
		rootFS.DiffIDs = nil

		if expected, actual := len(m.Layers), len(img.RootFS.DiffIDs); expected != actual {
			return fmt.Errorf("invalid manifest, layers length mismatch: expected %q, got %q", expected, actual)
		}

		for i, diffID := range img.RootFS.DiffIDs {
			layerPath, err := safePath(tmpDir, m.Layers[i])
			if err != nil {
				return err
			}
			r := rootFS
			r.Append(diffID)
			newLayer, err := l.ls.Get(r.ChainID())
			if err != nil {
				newLayer, err = l.loadLayer(layerPath, rootFS, diffID.String(), m.LayerSources[diffID], progressOutput)
				if err != nil {
					return err
				}
			}
			defer layer.ReleaseAndLog(l.ls, newLayer)
			if expected, actual := diffID, newLayer.DiffID(); expected != actual {
				return fmt.Errorf("invalid diffID for layer %d: expected %q, got %q", i, expected, actual)
			}
			rootFS.Append(diffID)
		}

		imgID, err := l.is.Create(config)
		if err != nil {
			return err
		}
		imageIDsStr += fmt.Sprintf("Loaded image ID: %s\n", imgID)

		imageRefCount = 0
		for _, repoTag := range m.RepoTags {
			named, err := reference.ParseNamed(repoTag)
			if err != nil {
				return err
			}
			ref, ok := named.(reference.NamedTagged)
			if !ok {
				return fmt.Errorf("invalid tag %q", repoTag)
			}
			l.setLoadedTag(ref, imgID, outStream)
			outStream.Write([]byte(fmt.Sprintf("Loaded image: %s\n", ref)))
			imageRefCount++
		}

		parentLinks = append(parentLinks, parentLink{imgID, m.Parent})
		l.loggerImgEvent.LogImageEvent(imgID.String(), imgID.String(), "load")
	}

	for _, p := range validatedParentLinks(parentLinks) {
		if p.parentID != "" {
			if err := l.setParentID(p.id, p.parentID); err != nil {
				return err
			}
		}
	}

	if imageRefCount == 0 {
		outStream.Write([]byte(imageIDsStr))
	}

	return nil
}

func (l *tarexporter) setParentID(id, parentID image.ID) error {
	img, err := l.is.Get(id)
	if err != nil {
		return err
	}
	parent, err := l.is.Get(parentID)
	if err != nil {
		return err
	}
	if !checkValidParent(img, parent) {
		return fmt.Errorf("image %v is not a valid parent for %v", parent.ID(), img.ID())
	}
	return l.is.SetParent(id, parentID)
}

func (l *tarexporter) loadLayer(filename string, rootFS image.RootFS, id string, foreignSrc distribution.Descriptor, progressOutput progress.Output) (layer.Layer, error) {
	rawTar, err := os.Open(filename)
	if err != nil {
		logrus.Debugf("Error reading embedded tar: %v", err)
		return nil, err
	}
	defer rawTar.Close()

	var r io.Reader
	if progressOutput != nil {
		fileInfo, err := rawTar.Stat()
		if err != nil {
			logrus.Debugf("Error statting file: %v", err)
			return nil, err
		}

		r = progress.NewProgressReader(rawTar, progressOutput, fileInfo.Size(), stringid.TruncateID(id), "Loading layer")
	} else {
		r = rawTar
	}

	inflatedLayerData, err := archive.DecompressStream(r)
	if err != nil {
		return nil, err
	}
	defer inflatedLayerData.Close()

	if ds, ok := l.ls.(layer.DescribableStore); ok {
		return ds.RegisterWithDescriptor(inflatedLayerData, rootFS.ChainID(), foreignSrc)
	}
	return l.ls.Register(inflatedLayerData, rootFS.ChainID())
}

func (l *tarexporter) setLoadedTag(ref reference.NamedTagged, imgID image.ID, outStream io.Writer) error {
	if prevID, err := l.rs.Get(ref); err == nil && prevID != imgID {
		fmt.Fprintf(outStream, "The image %s already exists, renaming the old one with ID %s to empty string\n", ref.String(), string(prevID)) // todo: this message is wrong in case of multiple tags
	}

	if err := l.rs.AddTag(ref, imgID, true); err != nil {
		return err
	}
	return nil
}

func (l *tarexporter) legacyLoad(tmpDir string, outStream io.Writer, progressOutput progress.Output) error {
	legacyLoadedMap := make(map[string]image.ID)

	dirs, err := ioutil.ReadDir(tmpDir)
	if err != nil {
		return err
	}

	// every dir represents an image
	for _, d := range dirs {
		if d.IsDir() {
			if err := l.legacyLoadImage(d.Name(), tmpDir, legacyLoadedMap, progressOutput); err != nil {
				return err
			}
		}
	}

	// load tags from repositories file
	repositoriesPath, err := safePath(tmpDir, legacyRepositoriesFileName)
	if err != nil {
		return err
	}
	repositoriesFile, err := os.Open(repositoriesPath)
	if err != nil {
		if !os.IsNotExist(err) {
			return err
		}
		return repositoriesFile.Close()
	}
	defer repositoriesFile.Close()

	repositories := make(map[string]map[string]string)
	if err := json.NewDecoder(repositoriesFile).Decode(&repositories); err != nil {
		return err
	}

	for name, tagMap := range repositories {
		for tag, oldID := range tagMap {
			imgID, ok := legacyLoadedMap[oldID]
			if !ok {
				return fmt.Errorf("invalid target ID: %v", oldID)
			}
			named, err := reference.WithName(name)
			if err != nil {
				return err
			}
			ref, err := reference.WithTag(named, tag)
			if err != nil {
				return err
			}
			l.setLoadedTag(ref, imgID, outStream)
		}
	}

	return nil
}

func (l *tarexporter) legacyLoadImage(oldID, sourceDir string, loadedMap map[string]image.ID, progressOutput progress.Output) error {
	if _, loaded := loadedMap[oldID]; loaded {
		return nil
	}
	configPath, err := safePath(sourceDir, filepath.Join(oldID, legacyConfigFileName))
	if err != nil {
		return err
	}
	imageJSON, err := ioutil.ReadFile(configPath)
	if err != nil {
		logrus.Debugf("Error reading json: %v", err)
		return err
	}

	var img struct{ Parent string }
	if err := json.Unmarshal(imageJSON, &img); err != nil {
		return err
	}

	var parentID image.ID
	if img.Parent != "" {
		for {
			var loaded bool
			if parentID, loaded = loadedMap[img.Parent]; !loaded {
				if err := l.legacyLoadImage(img.Parent, sourceDir, loadedMap, progressOutput); err != nil {
					return err
				}
			} else {
				break
			}
		}
	}

	// todo: try to connect with migrate code
	rootFS := image.NewRootFS()
	var history []image.History

	if parentID != "" {
		parentImg, err := l.is.Get(parentID)
		if err != nil {
			return err
		}

		rootFS = parentImg.RootFS
		history = parentImg.History
	}

	layerPath, err := safePath(sourceDir, filepath.Join(oldID, legacyLayerFileName))
	if err != nil {
		return err
	}
	newLayer, err := l.loadLayer(layerPath, *rootFS, oldID, distribution.Descriptor{}, progressOutput)
	if err != nil {
		return err
	}
	rootFS.Append(newLayer.DiffID())

	h, err := v1.HistoryFromConfig(imageJSON, false)
	if err != nil {
		return err
	}
	history = append(history, h)

	config, err := v1.MakeConfigFromV1Config(imageJSON, rootFS, history)
	if err != nil {
		return err
	}
	imgID, err := l.is.Create(config)
	if err != nil {
		return err
	}

	metadata, err := l.ls.Release(newLayer)
	layer.LogReleaseMetadata(metadata)
	if err != nil {
		return err
	}

	if parentID != "" {
		if err := l.is.SetParent(imgID, parentID); err != nil {
			return err
		}
	}

	loadedMap[oldID] = imgID
	return nil
}

func safePath(base, path string) (string, error) {
	return symlink.FollowSymlinkInScope(filepath.Join(base, path), base)
}

type parentLink struct {
	id, parentID image.ID
}

func validatedParentLinks(pl []parentLink) (ret []parentLink) {
mainloop:
	for i, p := range pl {
		ret = append(ret, p)
		for _, p2 := range pl {
			if p2.id == p.parentID && p2.id != p.id {
				continue mainloop
			}
		}
		ret[i].parentID = ""
	}
	return
}

func checkValidParent(img, parent *image.Image) bool {
	if len(img.History) == 0 && len(parent.History) == 0 {
		return true // having history is not mandatory
	}
	if len(img.History)-len(parent.History) != 1 {
		return false
	}
	for i, h := range parent.History {
		if !reflect.DeepEqual(h, img.History[i]) {
			return false
		}
	}
	return true
}
<|endoftext|>


<commit_before>// mandelterm - a terminal based Mandelbrot set viewer
package main

import (
	"fmt"
	"github.com/nsf/termbox-go"
	"log"
	"math/cmplx"
	"time"
)

// Constants
const (
	// On my terminal 80 x 24 screen is 800 x 504 pixels
	// So each character is 10 x 21 pixels
	// aspect ratio of terminal characters
	aspect = 2.1

	// Fraction of the radius we pan on each keypress
	pan = 0.2

	// Factor we zoom in on each keypress
	zoom = 2
)

// Globals
var (
	showHelp = true
	showInfo = true
	center   complex128
	radius   float64
	depth    int
)

// Reset to the start position
func Reset() {
	center = complex(0, 0)
	radius = 2.0
	depth = 256
}

// Print a string
func Print(x, y int, fg, bg termbox.Attribute, msg string) {
	for _, c := range msg {
		termbox.SetCell(x, y, c, fg, bg)
		x++
	}
}

// Printf a string
func Printf(x, y int, fg, bg termbox.Attribute, format string, args ...interface{}) {
	s := fmt.Sprintf(format, args...)
	Print(x, y, fg, bg, s)
}

// Help
func Help() {
	Print(5, 4, termbox.ColorRed, termbox.ColorWhite, "Terminal Mandlebrot by Nick Craig-Wood")
	Print(5, 6, termbox.ColorBlack, termbox.ColorWhite, "* Arrow keys to Move")
	Print(5, 7, termbox.ColorBlack, termbox.ColorWhite, "* PgUp/PdDown/+/- to zoom")
	Print(5, 8, termbox.ColorBlack, termbox.ColorWhite, "* [/] to change depth")
	Print(5, 9, termbox.ColorBlack, termbox.ColorWhite, "* h toggle help on and off")
	Print(5, 10, termbox.ColorBlack, termbox.ColorWhite, "* i toggle info on and off")
	Print(5, 11, termbox.ColorBlack, termbox.ColorWhite, "* q/ESC/c-C to quit")
	Print(5, 12, termbox.ColorBlack, termbox.ColorWhite, "* r to reset to start")
}

// Info
func Info(dt time.Duration) {
	Printf(0, 0, termbox.ColorBlack, termbox.ColorWhite, "c = %g, r = %g, Depth %d, rendered in %s", center, radius, depth, dt)
}

// Draw the mandelbrot set
func Draw() {
	start := time.Now()
	w, h := termbox.Size()
	// Choose shortest direction for reference
	var dx, dy float64
	if float64(h) > float64(w)/aspect {
		dx = 2 * radius / float64(w)
		dy = 2 * radius / float64(w) * aspect
	} else {
		dx = 2 * radius / float64(h) / aspect
		dy = 2 * radius / float64(h)
	}

	// FIXME aspect ratio
	termbox.Clear(termbox.ColorDefault, termbox.ColorDefault)
	fy := imag(center) + dy*float64(-h/2)
	for y := 0; y < h; y++ {
		fx := real(center) + dx*float64(-w/2)
		for x := 0; x < w; x++ {
			z := complex(0, 0)
			c := complex(fx, fy)
			var i int
			for i = 0; i < depth; i++ {
				if cmplx.Abs(z) >= 2 {
					break
				}
				z = z*z + c
			}
			termbox.SetCell(x, y, ' ', termbox.ColorDefault, termbox.Attribute(i%8)+1)
			fx += dx
		}
		fy += dy
	}
	if showInfo {
		Info(time.Now().Sub(start))
	}
	if showHelp {
		Help()
	}
	termbox.Flush()
}

func main() {
	err := termbox.Init()
	if err != nil {
		log.Fatal(err)
	}
	defer termbox.Close()
	Reset()
	Draw()
	for {
		ev := termbox.PollEvent()
		if ev.Type == termbox.EventKey {
			switch ev.Key + termbox.Key(ev.Ch) {
			case termbox.KeyEsc, termbox.KeyCtrlC, 'q':
				return
			case termbox.KeyArrowUp:
				center += complex(0.0, -radius*pan)
			case termbox.KeyArrowDown:
				center += complex(0.0, radius*pan)
			case termbox.KeyArrowLeft:
				center += complex(-radius*pan, 0.0)
			case termbox.KeyArrowRight:
				center += complex(radius*pan, 0.0)
			case termbox.KeyPgup, '=', '+':
				radius /= zoom
			case termbox.KeyPgdn, '-', '_':
				radius *= zoom
			case ']':
				depth *= 2
			case '[':
				depth /= 2
				if depth < 64 {
					depth = 64
				}
			case 'h':
				showHelp = !showHelp
			case 'i':
				showInfo = !showInfo
			case 'r':
				Reset()
			}
		}
		Draw()
	}
}
<commit_msg>Make it run in parallel using all CPUs<commit_after>// mandelterm - a terminal based Mandelbrot set viewer
package main

import (
	"fmt"
	"github.com/nsf/termbox-go"
	"log"
	"math/cmplx"
	"runtime"
	"sync"
	"time"
)

// Constants
const (
	// On my terminal 80 x 24 screen is 800 x 504 pixels
	// So each character is 10 x 21 pixels
	// aspect ratio of terminal characters
	aspect = 2.1

	// Fraction of the radius we pan on each keypress
	pan = 0.2

	// Factor we zoom in on each keypress
	zoom = 2
)

// Globals
var (
	showHelp = true
	showInfo = true
	center   complex128
	radius   float64
	depth    int
)

// Reset to the start position
func Reset() {
	center = complex(0, 0)
	radius = 2.0
	depth = 256
}

// Print a string
func Print(x, y int, fg, bg termbox.Attribute, msg string) {
	for _, c := range msg {
		termbox.SetCell(x, y, c, fg, bg)
		x++
	}
}

// Printf a string
func Printf(x, y int, fg, bg termbox.Attribute, format string, args ...interface{}) {
	s := fmt.Sprintf(format, args...)
	Print(x, y, fg, bg, s)
}

// Help
func Help() {
	Print(0, 0, termbox.ColorRed, termbox.ColorWhite, "Terminal Mandlebrot by Nick Craig-Wood")
	Print(0, 1, termbox.ColorBlack, termbox.ColorWhite, "* Arrow keys to Move")
	Print(0, 2, termbox.ColorBlack, termbox.ColorWhite, "* PgUp/PdDown/+/- to zoom")
	Print(0, 3, termbox.ColorBlack, termbox.ColorWhite, "* [/] to change depth")
	Print(0, 4, termbox.ColorBlack, termbox.ColorWhite, "* h toggle help on and off")
	Print(0, 5, termbox.ColorBlack, termbox.ColorWhite, "* i toggle info on and off")
	Print(0, 6, termbox.ColorBlack, termbox.ColorWhite, "* q/ESC/c-C to quit")
	Print(0, 7, termbox.ColorBlack, termbox.ColorWhite, "* r to reset to start")
}

// Info
func Info(dt time.Duration) {
	_, h := termbox.Size()
	Printf(0, h-1, termbox.ColorBlack, termbox.ColorWhite, "c = %g, r = %g, Depth %d, rendered in %s", center, radius, depth, dt)
}

// Plot a horizontal line from the mandelbrot set
func calculateLine(fx, fy, dx float64, line []termbox.Attribute, wg *sync.WaitGroup) {
	defer wg.Done()
	for x := range line {
		z := complex(0, 0)
		c := complex(fx, fy)
		var i int
		for i = 0; i < depth; i++ {
			if cmplx.Abs(z) >= 2 {
				break
			}
			z = z*z + c
		}
		line[x] = termbox.Attribute(i%8) + 1
		fx += dx
	}
}

// Draw the mandelbrot set
func Draw() {
	start := time.Now()
	w, h := termbox.Size()

	// Choose shortest direction for radius
	var dx, dy float64
	if float64(h) > float64(w)/aspect {
		dx = 2 * radius / float64(w)
		dy = 2 * radius / float64(w) * aspect
	} else {
		dx = 2 * radius / float64(h) / aspect
		dy = 2 * radius / float64(h)
	}

	// Calculate mandelbrot into screen
	screen := make([][]termbox.Attribute, h)
	fy := imag(center) + dy*float64(-h/2)
	var wg sync.WaitGroup
	for y := range screen {
		fx := real(center) + dx*float64(-w/2)
		screen[y] = make([]termbox.Attribute, w)
		wg.Add(1)
		go calculateLine(fx, fy, dx, screen[y], &wg)
		fy += dy

	}
	wg.Wait()

	// Plot
	termbox.Clear(termbox.ColorDefault, termbox.ColorDefault)
	for y := range screen {
		for x, col := range screen[y] {
			termbox.SetCell(x, y, ' ', termbox.ColorDefault, col)
		}
	}
	if showInfo {
		Info(time.Now().Sub(start))
	}
	if showHelp {
		Help()
	}
	termbox.Flush()
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	err := termbox.Init()
	if err != nil {
		log.Fatal(err)
	}
	defer termbox.Close()
	Reset()
	Draw()
	for {
		ev := termbox.PollEvent()
		if ev.Type == termbox.EventKey {
			switch ev.Key + termbox.Key(ev.Ch) {
			case termbox.KeyEsc, termbox.KeyCtrlC, 'q':
				return
			case termbox.KeyArrowUp:
				center += complex(0.0, -radius*pan)
			case termbox.KeyArrowDown:
				center += complex(0.0, radius*pan)
			case termbox.KeyArrowLeft:
				center += complex(-radius*pan, 0.0)
			case termbox.KeyArrowRight:
				center += complex(radius*pan, 0.0)
			case termbox.KeyPgup, '=', '+':
				radius /= zoom
			case termbox.KeyPgdn, '-', '_':
				radius *= zoom
			case ']':
				depth *= 2
			case '[':
				depth /= 2
				if depth < 64 {
					depth = 64
				}
			case 'h':
				showHelp = !showHelp
			case 'i':
				showInfo = !showInfo
			case 'r':
				Reset()
			}
		}
		Draw()
	}
}
<|endoftext|>


<commit_before>package copy

import (
	"io"

	internalblobinfocache "github.com/containers/image/v5/internal/blobinfocache"
	"github.com/containers/image/v5/pkg/compression"
	compressiontypes "github.com/containers/image/v5/pkg/compression/types"
	"github.com/containers/image/v5/types"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

// bpDetectCompressionStepData contains data that the copy pipeline needs about the “detect compression” step.
type bpDetectCompressionStepData struct {
	isCompressed bool
	format       compressiontypes.Algorithm        // Valid if isCompressed
	decompressor compressiontypes.DecompressorFunc // Valid if isCompressed
}

// blobPipelineDetectCompressionStep updates *stream to detect its current compression format.
// srcInfo is only used for error messages.
// Returns data for other steps.
func blobPipelineDetectCompressionStep(stream *sourceStream, srcInfo types.BlobInfo) (bpDetectCompressionStepData, error) {
	// This requires us to “peek ahead” into the stream to read the initial part, which requires us to chain through another io.Reader returned by DetectCompression.
	format, decompressor, reader, err := compression.DetectCompressionFormat(stream.reader) // We could skip this in some cases, but let's keep the code path uniform
	if err != nil {
		return bpDetectCompressionStepData{}, errors.Wrapf(err, "reading blob %s", srcInfo.Digest)
	}
	stream.reader = reader

	res := bpDetectCompressionStepData{
		isCompressed: decompressor != nil,
		format:       format,
		decompressor: decompressor,
	}
	if expectedFormat, known := expectedCompressionFormats[stream.info.MediaType]; known && res.isCompressed && format.Name() != expectedFormat.Name() {
		logrus.Debugf("blob %s with type %s should be compressed with %s, but compressor appears to be %s", srcInfo.Digest.String(), srcInfo.MediaType, expectedFormat.Name(), format.Name())
	}
	return res, nil
}

// bpCompressionStepData contains data that the copy pipeline needs about the compression step.
type bpCompressionStepData struct {
	operation            types.LayerCompression      // Operation to use for updating the blob metadata.
	uploadedAlgorithm    *compressiontypes.Algorithm // An algorithm parameter for the compressionOperation edits.
	compressionMetadata  map[string]string           // Annotations that should be set on the uploaded blob. WARNING: This is only set after the srcStream.reader is fully consumed.
	srcCompressorName    string                      // Compressor name to record in the blob info cache for the source blob.
	uploadCompressorName string                      // Compressor name to record in the blob info cache for the uploaded blob.
	closers              []io.Closer                 // Objects to close after the upload is done, if any.
}

// blobPipelineCompressionStep updates *stream to compress and/or decompress it.
// srcInfo is only used for error messages.
// Returns data for other steps; the caller should eventually call updateCompressionEdits and perhaps recordValidatedBlobData,
// and must eventually call close.
func (c *copier) blobPipelineCompressionStep(stream *sourceStream, canModifyBlob bool,
	detected bpDetectCompressionStepData) (*bpCompressionStepData, error) {
	// WARNING: If you are adding new reasons to change the blob, update also the OptimizeDestinationImageAlreadyExists
	// short-circuit conditions
	compressionMetadata := map[string]string{}
	var operation types.LayerCompression
	var uploadedAlgorithm *compressiontypes.Algorithm
	srcCompressorName := internalblobinfocache.Uncompressed
	if detected.isCompressed {
		srcCompressorName = detected.format.Name()
	}
	var uploadCompressorName string
	var closers []io.Closer
	succeeded := false
	defer func() {
		if !succeeded {
			for _, c := range closers {
				c.Close()
			}
		}
	}()
	if canModifyBlob && isOciEncrypted(stream.info.MediaType) {
		// PreserveOriginal due to any compression not being able to be done on an encrypted blob unless decrypted
		logrus.Debugf("Using original blob without modification for encrypted blob")
		operation = types.PreserveOriginal
		srcCompressorName = internalblobinfocache.UnknownCompression
		uploadedAlgorithm = nil
		uploadCompressorName = internalblobinfocache.UnknownCompression
	} else if canModifyBlob && c.dest.DesiredLayerCompression() == types.Compress && !detected.isCompressed {
		logrus.Debugf("Compressing blob on the fly")
		operation = types.Compress
		pipeReader, pipeWriter := io.Pipe()
		closers = append(closers, pipeReader)

		if c.compressionFormat != nil {
			uploadedAlgorithm = c.compressionFormat
		} else {
			uploadedAlgorithm = defaultCompressionFormat
		}
		// If this fails while writing data, it will do pipeWriter.CloseWithError(); if it fails otherwise,
		// e.g. because we have exited and due to pipeReader.Close() above further writing to the pipe has failed,
		// we don’t care.
		go c.compressGoroutine(pipeWriter, stream.reader, compressionMetadata, *uploadedAlgorithm) // Closes pipeWriter
		stream.reader = pipeReader
		stream.info = types.BlobInfo{ // FIXME? Should we preserve more data in src.info?
			Digest: "",
			Size:   -1,
		}
		uploadCompressorName = uploadedAlgorithm.Name()
	} else if canModifyBlob && c.dest.DesiredLayerCompression() == types.Compress && detected.isCompressed &&
		c.compressionFormat != nil && c.compressionFormat.Name() != detected.format.Name() {
		// When the blob is compressed, but the desired format is different, it first needs to be decompressed and finally
		// re-compressed using the desired format.
		logrus.Debugf("Blob will be converted")

		operation = types.PreserveOriginal
		s, err := detected.decompressor(stream.reader)
		if err != nil {
			return nil, err
		}
		closers = append(closers, s)

		pipeReader, pipeWriter := io.Pipe()
		closers = append(closers, pipeReader)

		uploadedAlgorithm = c.compressionFormat
		go c.compressGoroutine(pipeWriter, s, compressionMetadata, *uploadedAlgorithm) // Closes pipeWriter

		stream.reader = pipeReader
		stream.info = types.BlobInfo{ // FIXME? Should we preserve more data in src.info?
			Digest: "",
			Size:   -1,
		}
		uploadCompressorName = uploadedAlgorithm.Name()
	} else if canModifyBlob && c.dest.DesiredLayerCompression() == types.Decompress && detected.isCompressed {
		logrus.Debugf("Blob will be decompressed")
		operation = types.Decompress
		s, err := detected.decompressor(stream.reader)
		if err != nil {
			return nil, err
		}
		closers = append(closers, s)
		stream.reader = s
		stream.info = types.BlobInfo{ // FIXME? Should we preserve more data in src.info?
			Digest: "",
			Size:   -1,
		}
		uploadedAlgorithm = nil
		uploadCompressorName = internalblobinfocache.Uncompressed
	} else {
		// PreserveOriginal might also need to recompress the original blob if the desired compression format is different.
		logrus.Debugf("Using original blob without modification")
		operation = types.PreserveOriginal
		// Remember if the original blob was compressed, and if so how, so that if
		// LayerInfosForCopy() returned something that differs from what was in the
		// source's manifest, and UpdatedImage() needs to call UpdateLayerInfos(),
		// it will be able to correctly derive the MediaType for the copied blob.
		if detected.isCompressed {
			uploadedAlgorithm = &detected.format
		} else {
			uploadedAlgorithm = nil
		}
		uploadCompressorName = srcCompressorName
	}
	succeeded = true
	return &bpCompressionStepData{
		operation:            operation,
		uploadedAlgorithm:    uploadedAlgorithm,
		compressionMetadata:  compressionMetadata,
		srcCompressorName:    srcCompressorName,
		uploadCompressorName: uploadCompressorName,
		closers:              closers,
	}, nil
}

// updateCompressionEdits sets *operation, *algorithm and updates *annotations, if necessary.
func (d *bpCompressionStepData) updateCompressionEdits(operation *types.LayerCompression, algorithm **compressiontypes.Algorithm, annotations *map[string]string) {
	*operation = d.operation
	// If we can modify the layer's blob, set the desired algorithm for it to be set in the manifest.
	*algorithm = d.uploadedAlgorithm
	if *annotations == nil {
		*annotations = map[string]string{}
	}
	for k, v := range d.compressionMetadata {
		(*annotations)[k] = v
	}
}

// recordValidatedBlobData updates b.blobInfoCache with data about the created uploadedInfo adnd the original srcInfo.
// This must ONLY be called if all data has been validated by OUR code, and is not comming from third parties.
func (d *bpCompressionStepData) recordValidatedDigestData(c *copier, uploadedInfo types.BlobInfo, srcInfo types.BlobInfo,
	encryptionStep *bpEncryptionStepData, decryptionStep *bpDecryptionStepData) error {
	// Don’t record any associations that involve encrypted data. This is a bit crude,
	// some blob substitutions (replacing pulls of encrypted data with local reuse of known decryption outcomes)
	// might be safe, but it’s not trivially obvious, so let’s be conservative for now.
	// This crude approach also means we don’t need to record whether a blob is encrypted
	// in the blob info cache (which would probably be necessary for any more complex logic),
	// and the simplicity is attractive.
	if !encryptionStep.encrypting && !decryptionStep.decrypting {
		// If d.operation != types.PreserveOriginal, we now have two reliable digest values:
		// srcinfo.Digest describes the pre-d.operation input, verified by digestingReader
		// uploadedInfo.Digest describes the post-d.operation output, computed by PutBlob
		// (because stream.info.Digest == "", this must have been computed afresh).
		switch d.operation {
		case types.PreserveOriginal:
			break // Do nothing, we have only one digest and we might not have even verified it.
		case types.Compress:
			c.blobInfoCache.RecordDigestUncompressedPair(uploadedInfo.Digest, srcInfo.Digest)
		case types.Decompress:
			c.blobInfoCache.RecordDigestUncompressedPair(srcInfo.Digest, uploadedInfo.Digest)
		default:
			return errors.Errorf("Internal error: Unexpected d.operation value %#v", d.operation)
		}
	}
	if d.uploadCompressorName != "" && d.uploadCompressorName != internalblobinfocache.UnknownCompression {
		c.blobInfoCache.RecordDigestCompressorName(uploadedInfo.Digest, d.uploadCompressorName)
	}
	if srcInfo.Digest != "" && d.srcCompressorName != "" && d.srcCompressorName != internalblobinfocache.UnknownCompression {
		c.blobInfoCache.RecordDigestCompressorName(srcInfo.Digest, d.srcCompressorName)
	}
	return nil
}

// close closes objects that carry state throughout the compression/decompression operation.
func (d *bpCompressionStepData) close() {
	for _, c := range d.closers {
		c.Close()
	}
}

// doCompression reads all input from src and writes its compressed equivalent to dest.
func doCompression(dest io.Writer, src io.Reader, metadata map[string]string, compressionFormat compressiontypes.Algorithm, compressionLevel *int) error {
	compressor, err := compression.CompressStreamWithMetadata(dest, metadata, compressionFormat, compressionLevel)
	if err != nil {
		return err
	}

	buf := make([]byte, compressionBufferSize)

	_, err = io.CopyBuffer(compressor, src, buf) // Sets err to nil, i.e. causes dest.Close()
	if err != nil {
		compressor.Close()
		return err
	}

	return compressor.Close()
}

// compressGoroutine reads all input from src and writes its compressed equivalent to dest.
func (c *copier) compressGoroutine(dest *io.PipeWriter, src io.Reader, metadata map[string]string, compressionFormat compressiontypes.Algorithm) {
	err := errors.New("Internal error: unexpected panic in compressGoroutine")
	defer func() { // Note that this is not the same as {defer dest.CloseWithError(err)}; we need err to be evaluated lazily.
		_ = dest.CloseWithError(err) // CloseWithError(nil) is equivalent to Close(), always returns nil
	}()

	err = doCompression(dest, src, metadata, compressionFormat, c.compressionLevel)
}
<commit_msg>Rename uploadCompressorName to uploadedCompressorName<commit_after>package copy

import (
	"io"

	internalblobinfocache "github.com/containers/image/v5/internal/blobinfocache"
	"github.com/containers/image/v5/pkg/compression"
	compressiontypes "github.com/containers/image/v5/pkg/compression/types"
	"github.com/containers/image/v5/types"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

// bpDetectCompressionStepData contains data that the copy pipeline needs about the “detect compression” step.
type bpDetectCompressionStepData struct {
	isCompressed bool
	format       compressiontypes.Algorithm        // Valid if isCompressed
	decompressor compressiontypes.DecompressorFunc // Valid if isCompressed
}

// blobPipelineDetectCompressionStep updates *stream to detect its current compression format.
// srcInfo is only used for error messages.
// Returns data for other steps.
func blobPipelineDetectCompressionStep(stream *sourceStream, srcInfo types.BlobInfo) (bpDetectCompressionStepData, error) {
	// This requires us to “peek ahead” into the stream to read the initial part, which requires us to chain through another io.Reader returned by DetectCompression.
	format, decompressor, reader, err := compression.DetectCompressionFormat(stream.reader) // We could skip this in some cases, but let's keep the code path uniform
	if err != nil {
		return bpDetectCompressionStepData{}, errors.Wrapf(err, "reading blob %s", srcInfo.Digest)
	}
	stream.reader = reader

	res := bpDetectCompressionStepData{
		isCompressed: decompressor != nil,
		format:       format,
		decompressor: decompressor,
	}
	if expectedFormat, known := expectedCompressionFormats[stream.info.MediaType]; known && res.isCompressed && format.Name() != expectedFormat.Name() {
		logrus.Debugf("blob %s with type %s should be compressed with %s, but compressor appears to be %s", srcInfo.Digest.String(), srcInfo.MediaType, expectedFormat.Name(), format.Name())
	}
	return res, nil
}

// bpCompressionStepData contains data that the copy pipeline needs about the compression step.
type bpCompressionStepData struct {
	operation              types.LayerCompression      // Operation to use for updating the blob metadata.
	uploadedAlgorithm      *compressiontypes.Algorithm // An algorithm parameter for the compressionOperation edits.
	compressionMetadata    map[string]string           // Annotations that should be set on the uploaded blob. WARNING: This is only set after the srcStream.reader is fully consumed.
	srcCompressorName      string                      // Compressor name to record in the blob info cache for the source blob.
	uploadedCompressorName string                      // Compressor name to record in the blob info cache for the uploaded blob.
	closers                []io.Closer                 // Objects to close after the upload is done, if any.
}

// blobPipelineCompressionStep updates *stream to compress and/or decompress it.
// srcInfo is only used for error messages.
// Returns data for other steps; the caller should eventually call updateCompressionEdits and perhaps recordValidatedBlobData,
// and must eventually call close.
func (c *copier) blobPipelineCompressionStep(stream *sourceStream, canModifyBlob bool,
	detected bpDetectCompressionStepData) (*bpCompressionStepData, error) {
	// WARNING: If you are adding new reasons to change the blob, update also the OptimizeDestinationImageAlreadyExists
	// short-circuit conditions
	compressionMetadata := map[string]string{}
	var operation types.LayerCompression
	var uploadedAlgorithm *compressiontypes.Algorithm
	srcCompressorName := internalblobinfocache.Uncompressed
	if detected.isCompressed {
		srcCompressorName = detected.format.Name()
	}
	var uploadedCompressorName string
	var closers []io.Closer
	succeeded := false
	defer func() {
		if !succeeded {
			for _, c := range closers {
				c.Close()
			}
		}
	}()
	if canModifyBlob && isOciEncrypted(stream.info.MediaType) {
		// PreserveOriginal due to any compression not being able to be done on an encrypted blob unless decrypted
		logrus.Debugf("Using original blob without modification for encrypted blob")
		operation = types.PreserveOriginal
		srcCompressorName = internalblobinfocache.UnknownCompression
		uploadedAlgorithm = nil
		uploadedCompressorName = internalblobinfocache.UnknownCompression
	} else if canModifyBlob && c.dest.DesiredLayerCompression() == types.Compress && !detected.isCompressed {
		logrus.Debugf("Compressing blob on the fly")
		operation = types.Compress
		pipeReader, pipeWriter := io.Pipe()
		closers = append(closers, pipeReader)

		if c.compressionFormat != nil {
			uploadedAlgorithm = c.compressionFormat
		} else {
			uploadedAlgorithm = defaultCompressionFormat
		}
		// If this fails while writing data, it will do pipeWriter.CloseWithError(); if it fails otherwise,
		// e.g. because we have exited and due to pipeReader.Close() above further writing to the pipe has failed,
		// we don’t care.
		go c.compressGoroutine(pipeWriter, stream.reader, compressionMetadata, *uploadedAlgorithm) // Closes pipeWriter
		stream.reader = pipeReader
		stream.info = types.BlobInfo{ // FIXME? Should we preserve more data in src.info?
			Digest: "",
			Size:   -1,
		}
		uploadedCompressorName = uploadedAlgorithm.Name()
	} else if canModifyBlob && c.dest.DesiredLayerCompression() == types.Compress && detected.isCompressed &&
		c.compressionFormat != nil && c.compressionFormat.Name() != detected.format.Name() {
		// When the blob is compressed, but the desired format is different, it first needs to be decompressed and finally
		// re-compressed using the desired format.
		logrus.Debugf("Blob will be converted")

		operation = types.PreserveOriginal
		s, err := detected.decompressor(stream.reader)
		if err != nil {
			return nil, err
		}
		closers = append(closers, s)

		pipeReader, pipeWriter := io.Pipe()
		closers = append(closers, pipeReader)

		uploadedAlgorithm = c.compressionFormat
		go c.compressGoroutine(pipeWriter, s, compressionMetadata, *uploadedAlgorithm) // Closes pipeWriter

		stream.reader = pipeReader
		stream.info = types.BlobInfo{ // FIXME? Should we preserve more data in src.info?
			Digest: "",
			Size:   -1,
		}
		uploadedCompressorName = uploadedAlgorithm.Name()
	} else if canModifyBlob && c.dest.DesiredLayerCompression() == types.Decompress && detected.isCompressed {
		logrus.Debugf("Blob will be decompressed")
		operation = types.Decompress
		s, err := detected.decompressor(stream.reader)
		if err != nil {
			return nil, err
		}
		closers = append(closers, s)
		stream.reader = s
		stream.info = types.BlobInfo{ // FIXME? Should we preserve more data in src.info?
			Digest: "",
			Size:   -1,
		}
		uploadedAlgorithm = nil
		uploadedCompressorName = internalblobinfocache.Uncompressed
	} else {
		// PreserveOriginal might also need to recompress the original blob if the desired compression format is different.
		logrus.Debugf("Using original blob without modification")
		operation = types.PreserveOriginal
		// Remember if the original blob was compressed, and if so how, so that if
		// LayerInfosForCopy() returned something that differs from what was in the
		// source's manifest, and UpdatedImage() needs to call UpdateLayerInfos(),
		// it will be able to correctly derive the MediaType for the copied blob.
		if detected.isCompressed {
			uploadedAlgorithm = &detected.format
		} else {
			uploadedAlgorithm = nil
		}
		uploadedCompressorName = srcCompressorName
	}
	succeeded = true
	return &bpCompressionStepData{
		operation:              operation,
		uploadedAlgorithm:      uploadedAlgorithm,
		compressionMetadata:    compressionMetadata,
		srcCompressorName:      srcCompressorName,
		uploadedCompressorName: uploadedCompressorName,
		closers:                closers,
	}, nil
}

// updateCompressionEdits sets *operation, *algorithm and updates *annotations, if necessary.
func (d *bpCompressionStepData) updateCompressionEdits(operation *types.LayerCompression, algorithm **compressiontypes.Algorithm, annotations *map[string]string) {
	*operation = d.operation
	// If we can modify the layer's blob, set the desired algorithm for it to be set in the manifest.
	*algorithm = d.uploadedAlgorithm
	if *annotations == nil {
		*annotations = map[string]string{}
	}
	for k, v := range d.compressionMetadata {
		(*annotations)[k] = v
	}
}

// recordValidatedBlobData updates b.blobInfoCache with data about the created uploadedInfo adnd the original srcInfo.
// This must ONLY be called if all data has been validated by OUR code, and is not comming from third parties.
func (d *bpCompressionStepData) recordValidatedDigestData(c *copier, uploadedInfo types.BlobInfo, srcInfo types.BlobInfo,
	encryptionStep *bpEncryptionStepData, decryptionStep *bpDecryptionStepData) error {
	// Don’t record any associations that involve encrypted data. This is a bit crude,
	// some blob substitutions (replacing pulls of encrypted data with local reuse of known decryption outcomes)
	// might be safe, but it’s not trivially obvious, so let’s be conservative for now.
	// This crude approach also means we don’t need to record whether a blob is encrypted
	// in the blob info cache (which would probably be necessary for any more complex logic),
	// and the simplicity is attractive.
	if !encryptionStep.encrypting && !decryptionStep.decrypting {
		// If d.operation != types.PreserveOriginal, we now have two reliable digest values:
		// srcinfo.Digest describes the pre-d.operation input, verified by digestingReader
		// uploadedInfo.Digest describes the post-d.operation output, computed by PutBlob
		// (because stream.info.Digest == "", this must have been computed afresh).
		switch d.operation {
		case types.PreserveOriginal:
			break // Do nothing, we have only one digest and we might not have even verified it.
		case types.Compress:
			c.blobInfoCache.RecordDigestUncompressedPair(uploadedInfo.Digest, srcInfo.Digest)
		case types.Decompress:
			c.blobInfoCache.RecordDigestUncompressedPair(srcInfo.Digest, uploadedInfo.Digest)
		default:
			return errors.Errorf("Internal error: Unexpected d.operation value %#v", d.operation)
		}
	}
	if d.uploadedCompressorName != "" && d.uploadedCompressorName != internalblobinfocache.UnknownCompression {
		c.blobInfoCache.RecordDigestCompressorName(uploadedInfo.Digest, d.uploadedCompressorName)
	}
	if srcInfo.Digest != "" && d.srcCompressorName != "" && d.srcCompressorName != internalblobinfocache.UnknownCompression {
		c.blobInfoCache.RecordDigestCompressorName(srcInfo.Digest, d.srcCompressorName)
	}
	return nil
}

// close closes objects that carry state throughout the compression/decompression operation.
func (d *bpCompressionStepData) close() {
	for _, c := range d.closers {
		c.Close()
	}
}

// doCompression reads all input from src and writes its compressed equivalent to dest.
func doCompression(dest io.Writer, src io.Reader, metadata map[string]string, compressionFormat compressiontypes.Algorithm, compressionLevel *int) error {
	compressor, err := compression.CompressStreamWithMetadata(dest, metadata, compressionFormat, compressionLevel)
	if err != nil {
		return err
	}

	buf := make([]byte, compressionBufferSize)

	_, err = io.CopyBuffer(compressor, src, buf) // Sets err to nil, i.e. causes dest.Close()
	if err != nil {
		compressor.Close()
		return err
	}

	return compressor.Close()
}

// compressGoroutine reads all input from src and writes its compressed equivalent to dest.
func (c *copier) compressGoroutine(dest *io.PipeWriter, src io.Reader, metadata map[string]string, compressionFormat compressiontypes.Algorithm) {
	err := errors.New("Internal error: unexpected panic in compressGoroutine")
	defer func() { // Note that this is not the same as {defer dest.CloseWithError(err)}; we need err to be evaluated lazily.
		_ = dest.CloseWithError(err) // CloseWithError(nil) is equivalent to Close(), always returns nil
	}()

	err = doCompression(dest, src, metadata, compressionFormat, c.compressionLevel)
}
<|endoftext|>


<commit_before>package plugins

import (
	"fmt"
	"strings"

	"github.com/Sirupsen/logrus"
	"github.com/satori/go.uuid"

	"github.com/belak/go-seabird"
	"github.com/go-irc/irc"
)

func init() {
	seabird.RegisterPlugin("channel_track", newChannelTracker)
}

// Channel is an internal type for representing a channel.
type Channel struct {
	users map[string]bool
}

// HasUser returns true if the user is in the channel, otherwise
// false.
func (c *Channel) HasUser(user string) bool {
	return c.users[user]
}

// User is an type for representing a user.
type User struct {
	channels map[string]bool

	Nick string
	UUID string
}

// Channels returns which channels the user is currently in.
func (u *User) Channels() []string {
	var ret []string
	for k := range u.channels {
		ret = append(ret, k)
	}
	return ret
}

// InChannel returns true if the user is in the channel, otherwise
// false.
func (u *User) InChannel(channel string) bool {
	return u.channels[channel]
}

// ChannelTracker is a simple plugin which is only meant to track what
// channels the bot is in, and what users are in a channel. It also
// provides a uuid mapping to a user, so if a user's nick changes,
// we'll still have a sort of "session" to keep track of them.
type ChannelTracker struct {
	isupport *ISupportPlugin

	// Notes for internal fields. Be very careful when modifying the
	// state. Because we control all of this, it is valid to make the
	// assumption that if a user is in p.uuids, it will be possible to
	// find them in p.users.

	// Channels can't be renamed, so it's just a mapping of name to
	// channel object.
	channels map[string]*Channel

	// Users can be renamed so we key them on uuid. There's also a
	// separate nick to uuid mapping.
	users map[string]*User

	// This simply maps the nick to the uuid
	uuids map[string]string

	// Session cleanup callbacks
	cleanupCallbacks []func(u *User)
}

func newChannelTracker(bm *seabird.BasicMux, isupport *ISupportPlugin) *ChannelTracker {
	p := &ChannelTracker{
		isupport: isupport,
		channels: make(map[string]*Channel),
		users:    make(map[string]*User),
		uuids:    make(map[string]string),
	}

	bm.Event("JOIN", p.joinCallback)
	bm.Event("PART", p.partCallback)
	bm.Event("KICK", p.kickCallback)
	bm.Event("QUIT", p.quitCallback)
	bm.Event("NICK", p.nickCallback)

	bm.Event("353", p.namesCallback)
	// We don't need this for anything currently, so it's being
	// disabled but left around.
	//
	// bm.Event("366", p.endOfNamesCallback)

	return p
}

// Public interfaces

// LookupUser will return the User object for the given nick or nil if
// we don't know about this user. The returned value can be stored and
// will track this user even if they change nicks.
func (p *ChannelTracker) LookupUser(user string) *User {
	userUUID, ok := p.uuids[user]
	if !ok {
		return nil
	}

	return p.users[userUUID]
}

// UsersInChannel will return all the users in the given channel name
// or nil if we're not in that channel.
func (p *ChannelTracker) UsersInChannel(channel string) []*User {
	c := p.LookupChannel(channel)
	if c == nil {
		return nil
	}

	var ret []*User
	for userUUID := range c.users {
		ret = append(ret, p.users[userUUID])
	}
	return ret
}

// LookupChannel will return the Channel object for the given channel
// name or nil if we're not in that channel.
func (p *ChannelTracker) LookupChannel(channel string) *Channel {
	return p.channels[channel]
}

// Channels will return all the channel objects this bot knows about.
func (p *ChannelTracker) Channels() []*Channel {
	var ret []*Channel
	for _, v := range p.channels {
		ret = append(ret, v)
	}
	return ret
}

// RegisterSessionCleanupCallback lets you register a function to be
// called when a session is removed.
func (p *ChannelTracker) RegisterSessionCleanupCallback(f func(u *User)) {
	p.cleanupCallbacks = append(p.cleanupCallbacks, f)
}

// Private functions

func (p *ChannelTracker) joinCallback(b *seabird.Bot, m *irc.Message) {
	user := m.Prefix.Name
	channel := m.Trailing()

	p.addUserToChannel(b, user, channel)

	//fmt.Printf("%s (%s) joined %s\n", user, p.uuids[user], channel)
}

func (p *ChannelTracker) partCallback(b *seabird.Bot, m *irc.Message) {
	user := m.Prefix.Name
	channel := m.Params[0]

	p.removeUserFromChannel(b, user, channel)

	//fmt.Printf("%s (%s) left %s\n", user, p.uuids[user], channel)
}

func (p *ChannelTracker) kickCallback(b *seabird.Bot, m *irc.Message) {
	//actor := m.Prefix.Name
	user := m.Params[1]
	channel := m.Params[0]

	p.removeUserFromChannel(b, user, channel)

	//fmt.Printf("%s (%s) kicked %s (%s) from %s\n", actor, p.uuids[actor], user, p.uuids[user], channel)
}

func (p *ChannelTracker) quitCallback(b *seabird.Bot, m *irc.Message) {
	user := m.Prefix.Name

	p.removeUser(b, user)

	//fmt.Printf("%s (%s) quit\n", user, p.uuids[user])
}

func (p *ChannelTracker) nickCallback(b *seabird.Bot, m *irc.Message) {
	oldUser := m.Prefix.Name
	newUser := m.Params[0]

	p.renameUser(b, oldUser, newUser)

	//fmt.Printf("%s (%s) changed their name to %s\n", oldUser, p.uuids[newUser], newUser)
}

func (p *ChannelTracker) namesCallback(b *seabird.Bot, m *irc.Message) {
	logger := b.GetLogger()

	// Sample: (qaohv)~&@%+
	prefix, _ := p.isupport.GetRaw("PREFIX")

	// We only care about the symbols
	i := strings.IndexByte(prefix, ')')
	if len(prefix) == 0 || prefix[0] != '(' || i < 0 {
		logger.WithField("prefix", prefix).Warnf("Invalid prefix format")
		return
	}

	prefixes := prefix[i:]

	channel := m.Params[2]
	users := strings.Split(m.Trailing(), " ")
	for _, user := range users {
		user = strings.TrimLeft(user, prefixes)

		// The bot user should be added via JOIN
		if user == b.CurrentNick() {
			continue
		}

		p.addUserToChannel(b, user, channel)

		//fmt.Printf("%s (%s) is in channel %s\n", user, p.uuids[user], channel)
	}
}

func (p *ChannelTracker) endOfNamesCallback(b *seabird.Bot, m *irc.Message) {
	channel := m.Params[1]

	fmt.Printf("Got all names for %s\n", channel)
}

// Implementation below

func (p *ChannelTracker) addUserToChannel(b *seabird.Bot, user, channel string) {
	logger := b.GetLogger().WithFields(logrus.Fields{
		"channel": channel,
		"user":    user,
	})

	// If the current user is joining a channel, we need to add it
	// before adding our user.
	if user == b.CurrentNick() {
		p.addChannel(b, channel)
	}

	// If we're not in this channel, issue a warning and bail.
	c, ok := p.channels[channel]
	if !ok {
		logger.Warn("Error adding user: bot not in channel")
		return
	}

	u := p.LookupUser(user)
	if u == nil {
		u = &User{
			Nick:     user,
			UUID:     uuid.NewV4().String(),
			channels: make(map[string]bool),
		}
		p.users[u.UUID] = u
		p.uuids[user] = u.UUID
	}

	logger = logger.WithField("userUUID", u.UUID)

	if _, ok := u.channels[channel]; ok {
		logger.Warn("User already in channel")
		return
	}

	u.channels[channel] = true
	c.users[u.UUID] = true

	logger.Info("User added to channel")
}

func (p *ChannelTracker) removeUserFromChannel(b *seabird.Bot, user, channel string) {
	logger := b.GetLogger().WithField("channel", channel)

	if user == b.CurrentNick() {
		p.removeChannel(b, channel)
	} else {
		u := p.LookupUser(user)
		if u == nil {
			logger.Warn("Can't remove unknown user")
			return
		}

		logger = logger.WithField("userUUID", u.UUID)

		if _, ok := u.channels[channel]; !ok {
			logger.Warn("Can only remove users from users they are in")
			return
		}

		delete(u.channels, channel)

		logger.Info("Removing user from channel")

		if len(u.channels) == 0 {
			p.removeUser(b, user)
		}

		logger.Info("Removed user from channel")
	}
}

func (p *ChannelTracker) addChannel(b *seabird.Bot, channel string) {
	logger := b.GetLogger().WithField("channel", channel)

	if _, ok := p.channels[channel]; ok {
		logger.Warn("Already in channel")
		return
	}

	p.channels[channel] = &Channel{
		users: make(map[string]bool),
	}

	logger.Info("Added channel")
}

func (p *ChannelTracker) removeChannel(b *seabird.Bot, channel string) {
	logger := b.GetLogger().WithField("channel", channel)

	c, ok := p.channels[channel]
	if !ok {
		logger.Warn("Can only remove channels we are in")
		return
	}

	// We only do this for channels because it can have side effects
	logger.Info("Removing channel")

	// Remove all users currently in this channel from this channel.
	for userUUID := range c.users {
		u := p.users[userUUID]
		delete(u.channels, channel)

		// If this user has no more channels, they need to be removed.
		if len(u.channels) == 0 {
			p.removeUser(b, u.Nick)
		}
	}

	// Clean up the channel object
	c.users = make(map[string]bool)

	// Remove the channel from tracking
	delete(p.channels, channel)

	logger.Info("Removed channel")
}

func (p *ChannelTracker) removeUser(b *seabird.Bot, user string) {
	logger := b.GetLogger().WithField("user", user)

	u := p.LookupUser(user)
	if u == nil {
		logger.Warn("User does not exist")
		return
	}

	// We need to clear out the channels and Nick to show this session
	// is invalid.
	u.Nick = ""
	for channel := range u.channels {
		delete(p.channels[channel].users, u.UUID)
	}
	u.channels = make(map[string]bool)

	// Now that the User is empty, delete all internal traces.
	delete(p.uuids, user)
	delete(p.users, u.UUID)

	// Run any cleanup callbacks
	for _, f := range p.cleanupCallbacks {
		f(u)
	}

	logger.Info("Removed user")
}

func (p *ChannelTracker) renameUser(b *seabird.Bot, oldNick, newNick string) {
	logger := b.GetLogger().WithFields(logrus.Fields{
		"oldNick": oldNick,
		"newNick": newNick,
	})

	u := p.LookupUser(oldNick)
	if u == nil {
		logger.Warn("Can't rename user that doesn't exist")
		return
	}

	logger = logger.WithField("userUUID", u.UUID)

	// Rename the user object
	u.Nick = newNick

	// Swap where the UUID points to
	delete(p.uuids, oldNick)
	p.uuids[newNick] = u.UUID

	logger.Info("Renamed user")
}
<commit_msg>Initial attempt to track +o, +v, etc<commit_after>package plugins

import (
	"fmt"
	"strings"

	"github.com/Sirupsen/logrus"
	"github.com/satori/go.uuid"

	"github.com/belak/go-seabird"
	"github.com/go-irc/irc"
)

func init() {
	seabird.RegisterPlugin("channel_track", newChannelTracker)
}

// Channel is an internal type for representing a channel.
type Channel struct {
	users map[string]bool
}

// HasUser returns true if the user is in the channel, otherwise
// false.
func (c *Channel) HasUser(user string) bool {
	return c.users[user]
}

// User is an type for representing a user.
type User struct {
	channels map[string]map[rune]bool
	Nick     string
	UUID     string
}

// Channels returns which channels the user is currently in.
func (u *User) Channels() []string {
	var ret []string
	for k := range u.channels {
		ret = append(ret, k)
	}
	return ret
}

// ModesInChannel returns a mapping of channel modes to a bool indicating if
// it's on or not for this user in this channel.
func (u *User) ModesInChannel(channel string) map[rune]bool {
	ret, ok := u.channels[channel]
	if !ok {
		ret = make(map[rune]bool)
	}
	return ret
}

// InChannel returns true if the user is in the channel, otherwise
// false.
func (u *User) InChannel(channel string) bool {
	_, ok := u.channels[channel]
	return ok
}

// ChannelTracker is a simple plugin which is only meant to track what
// channels the bot is in, and what users are in a channel. It also
// provides a uuid mapping to a user, so if a user's nick changes,
// we'll still have a sort of "session" to keep track of them.
type ChannelTracker struct {
	isupport *ISupportPlugin

	// Notes for internal fields. Be very careful when modifying the
	// state. Because we control all of this, it is valid to make the
	// assumption that if a user is in p.uuids, it will be possible to
	// find them in p.users.

	// Channels can't be renamed, so it's just a mapping of name to
	// channel object.
	channels map[string]*Channel

	// Users can be renamed so we key them on uuid. There's also a
	// separate nick to uuid mapping.
	users map[string]*User

	// This simply maps the nick to the uuid
	uuids map[string]string

	// Session cleanup callbacks
	cleanupCallbacks []func(u *User)
}

func newChannelTracker(bm *seabird.BasicMux, isupport *ISupportPlugin) *ChannelTracker {
	p := &ChannelTracker{
		isupport: isupport,
		channels: make(map[string]*Channel),
		users:    make(map[string]*User),
		uuids:    make(map[string]string),
	}

	bm.Event("JOIN", p.joinCallback)
	bm.Event("PART", p.partCallback)
	bm.Event("KICK", p.kickCallback)
	bm.Event("QUIT", p.quitCallback)
	bm.Event("NICK", p.nickCallback)
	bm.Event("MODE", p.modeCallback)

	bm.Event("352", p.whoCallback)
	bm.Event("353", p.namesCallback)
	// We don't need this for anything currently, so it's being
	// disabled but left around.
	//
	// bm.Event("366", p.endOfNamesCallback)

	return p
}

// Public interfaces

// LookupUser will return the User object for the given nick or nil if
// we don't know about this user. The returned value can be stored and
// will track this user even if they change nicks.
func (p *ChannelTracker) LookupUser(user string) *User {
	userUUID, ok := p.uuids[user]
	if !ok {
		return nil
	}

	return p.users[userUUID]
}

// UsersInChannel will return all the users in the given channel name
// or nil if we're not in that channel.
func (p *ChannelTracker) UsersInChannel(channel string) []*User {
	c := p.LookupChannel(channel)
	if c == nil {
		return nil
	}

	var ret []*User
	for userUUID := range c.users {
		ret = append(ret, p.users[userUUID])
	}
	return ret
}

// LookupChannel will return the Channel object for the given channel
// name or nil if we're not in that channel.
func (p *ChannelTracker) LookupChannel(channel string) *Channel {
	return p.channels[channel]
}

// Channels will return all the channel objects this bot knows about.
func (p *ChannelTracker) Channels() []*Channel {
	var ret []*Channel
	for _, v := range p.channels {
		ret = append(ret, v)
	}
	return ret
}

// RegisterSessionCleanupCallback lets you register a function to be
// called when a session is removed.
func (p *ChannelTracker) RegisterSessionCleanupCallback(f func(u *User)) {
	p.cleanupCallbacks = append(p.cleanupCallbacks, f)
}

// Private functions

func (p *ChannelTracker) joinCallback(b *seabird.Bot, m *irc.Message) {
	user := m.Prefix.Name
	channel := m.Trailing()

	p.addUserToChannel(b, user, channel)

	//fmt.Printf("%s (%s) joined %s\n", user, p.uuids[user], channel)
}

func (p *ChannelTracker) partCallback(b *seabird.Bot, m *irc.Message) {
	user := m.Prefix.Name
	channel := m.Params[0]

	p.removeUserFromChannel(b, user, channel)

	//fmt.Printf("%s (%s) left %s\n", user, p.uuids[user], channel)
}

func (p *ChannelTracker) kickCallback(b *seabird.Bot, m *irc.Message) {
	//actor := m.Prefix.Name
	user := m.Params[1]
	channel := m.Params[0]

	p.removeUserFromChannel(b, user, channel)

	//fmt.Printf("%s (%s) kicked %s (%s) from %s\n", actor, p.uuids[actor], user, p.uuids[user], channel)
}

func (p *ChannelTracker) quitCallback(b *seabird.Bot, m *irc.Message) {
	user := m.Prefix.Name

	p.removeUser(b, user)

	//fmt.Printf("%s (%s) quit\n", user, p.uuids[user])
}

func (p *ChannelTracker) nickCallback(b *seabird.Bot, m *irc.Message) {
	oldUser := m.Prefix.Name
	newUser := m.Params[0]

	p.renameUser(b, oldUser, newUser)

	//fmt.Printf("%s (%s) changed their name to %s\n", oldUser, p.uuids[newUser], newUser)
}

func (p *ChannelTracker) modeCallback(b *seabird.Bot, m *irc.Message) {
	// We only care about MODE messages where a specific user is changed.
	if len(m.Params) < 3 {
		return
	}

	channel := m.Params[0]
	target := m.Params[2]

	//change := m.Params[1]
	//op := change[0]
	//mode := change[1]

	// Ensure we know about this user and this channel
	u := p.LookupUser(target)
	c := p.LookupChannel(channel)
	if u == nil || c == nil {
		// TODO: Warning
		return
	}

	// Just send a WHO request and clear out the modes for this user because
	// mode parsing is hard.
	u.channels[channel] = make(map[rune]bool)
	b.Writef("WHO :%s", target)
}

func (p *ChannelTracker) whoCallback(b *seabird.Bot, m *irc.Message) {
	// Filter out broken messages
	if len(m.Params) < 7 {
		return
	}

	prefixes, ok := p.getSymbolToPrefixMapping(b)
	if !ok {
		return
	}

	var (
		channel = m.Params[0]
		//user    = m.Params[1]
		//host    = m.Params[2]
		//server  = m.Params[3]
		nick  = m.Params[4]
		modes = m.Params[5]
		//extra   = m.Params[6]
	)

	u := p.LookupUser(nick)
	c := p.LookupChannel(channel)
	if u == nil || c == nil {
		// TODO: Warning
		return
	}

	// Modes starts with H/G for here/gone, so we skip that because we don't
	// care too much about tracking it.
	userPrefixes := modes[1:]

	// Clear out the modes and reset them
	u.channels[channel] = make(map[rune]bool)
	for _, v := range userPrefixes {
		mode := prefixes[v]
		u.channels[channel][mode] = true
	}
}

func (p *ChannelTracker) getSymbolToPrefixMapping(b *seabird.Bot) (map[rune]rune, bool) {
	logger := b.GetLogger()

	// Sample: (qaohv)~&@%+
	prefix, _ := p.isupport.GetRaw("PREFIX")

	// We only care about the symbols
	i := strings.IndexByte(prefix, ')')
	if len(prefix) == 0 || prefix[0] != '(' || i < 0 {
		logger.WithField("prefix", prefix).Warnf("Invalid prefix format")
		return nil, false
	}

	// TODO: This might not be right if unicode is used in prefix symbols.
	prefixes := make(map[rune]rune)
	for k, p := range prefix[i:] {
		prefixes[p] = rune(prefix[k+1])
	}

	return prefixes, true
}

func (p *ChannelTracker) namesCallback(b *seabird.Bot, m *irc.Message) {
	prefixes, ok := p.getSymbolToPrefixMapping(b)
	if !ok {
		return
	}

	channel := m.Params[2]
	users := strings.Split(m.Trailing(), " ")
	for _, user := range users {
		i := strings.IndexFunc(user, func(r rune) bool {
			_, ok := prefixes[r]
			return !ok
		})

		var userPrefixes string
		if i != -1 {
			user = user[i:]
			userPrefixes = user[:i]
		}

		// The bot user should be added via JOIN
		if user == b.CurrentNick() {
			continue
		}

		p.addUserToChannel(b, user, channel)

		u := p.LookupUser(user)
		if u == nil {
			continue
		}

		// Clear out the modes and reset them
		u.channels[channel] = make(map[rune]bool)
		for _, v := range userPrefixes {
			mode := prefixes[v]
			u.channels[channel][mode] = true
		}
		fmt.Printf("%+v, %+v, %+v\n", channel, user, u.channels[channel])

		//fmt.Printf("%s (%s) is in channel %s\n", user, p.uuids[user], channel)
	}
}

func (p *ChannelTracker) endOfNamesCallback(b *seabird.Bot, m *irc.Message) {
	channel := m.Params[1]

	fmt.Printf("Got all names for %s\n", channel)
}

// Implementation below

func (p *ChannelTracker) addUserToChannel(b *seabird.Bot, user, channel string) {
	logger := b.GetLogger().WithFields(logrus.Fields{
		"channel": channel,
		"user":    user,
	})

	// If the current user is joining a channel, we need to add it
	// before adding our user.
	if user == b.CurrentNick() {
		p.addChannel(b, channel)
	}

	// If we're not in this channel, issue a warning and bail.
	c, ok := p.channels[channel]
	if !ok {
		logger.Warn("Error adding user: bot not in channel")
		return
	}

	u := p.LookupUser(user)
	if u == nil {
		u = &User{
			Nick:     user,
			UUID:     uuid.NewV4().String(),
			channels: make(map[string]map[rune]bool),
		}
		p.users[u.UUID] = u
		p.uuids[user] = u.UUID
	}

	logger = logger.WithField("userUUID", u.UUID)

	if _, ok := u.channels[channel]; ok {
		logger.Warn("User already in channel")
		return
	}

	u.channels[channel] = make(map[rune]bool)
	c.users[u.UUID] = true

	logger.Info("User added to channel")
}

func (p *ChannelTracker) removeUserFromChannel(b *seabird.Bot, user, channel string) {
	logger := b.GetLogger().WithField("channel", channel)

	if user == b.CurrentNick() {
		p.removeChannel(b, channel)
	} else {
		u := p.LookupUser(user)
		if u == nil {
			logger.Warn("Can't remove unknown user")
			return
		}

		logger = logger.WithField("userUUID", u.UUID)

		if _, ok := u.channels[channel]; !ok {
			logger.Warn("Can only remove users from users they are in")
			return
		}

		delete(u.channels, channel)

		logger.Info("Removing user from channel")

		if len(u.channels) == 0 {
			p.removeUser(b, user)
		}

		logger.Info("Removed user from channel")
	}
}

func (p *ChannelTracker) addChannel(b *seabird.Bot, channel string) {
	logger := b.GetLogger().WithField("channel", channel)

	if _, ok := p.channels[channel]; ok {
		logger.Warn("Already in channel")
		return
	}

	p.channels[channel] = &Channel{
		users: make(map[string]bool),
	}

	logger.Info("Added channel")
}

func (p *ChannelTracker) removeChannel(b *seabird.Bot, channel string) {
	logger := b.GetLogger().WithField("channel", channel)

	c, ok := p.channels[channel]
	if !ok {
		logger.Warn("Can only remove channels we are in")
		return
	}

	// We only do this for channels because it can have side effects
	logger.Info("Removing channel")

	// Remove all users currently in this channel from this channel.
	for userUUID := range c.users {
		u := p.users[userUUID]
		delete(u.channels, channel)

		// If this user has no more channels, they need to be removed.
		if len(u.channels) == 0 {
			p.removeUser(b, u.Nick)
		}
	}

	// Clean up the channel object
	c.users = make(map[string]bool)

	// Remove the channel from tracking
	delete(p.channels, channel)

	logger.Info("Removed channel")
}

func (p *ChannelTracker) removeUser(b *seabird.Bot, user string) {
	logger := b.GetLogger().WithField("user", user)

	u := p.LookupUser(user)
	if u == nil {
		logger.Warn("User does not exist")
		return
	}

	// We need to clear out the channels and Nick to show this session
	// is invalid.
	u.Nick = ""
	for channel := range u.channels {
		delete(p.channels[channel].users, u.UUID)
	}
	u.channels = make(map[string]map[rune]bool)

	// Now that the User is empty, delete all internal traces.
	delete(p.uuids, user)
	delete(p.users, u.UUID)

	// Run any cleanup callbacks
	for _, f := range p.cleanupCallbacks {
		f(u)
	}

	logger.Info("Removed user")
}

func (p *ChannelTracker) renameUser(b *seabird.Bot, oldNick, newNick string) {
	logger := b.GetLogger().WithFields(logrus.Fields{
		"oldNick": oldNick,
		"newNick": newNick,
	})

	u := p.LookupUser(oldNick)
	if u == nil {
		logger.Warn("Can't rename user that doesn't exist")
		return
	}

	logger = logger.WithField("userUUID", u.UUID)

	// Rename the user object
	u.Nick = newNick

	// Swap where the UUID points to
	delete(p.uuids, oldNick)
	p.uuids[newNick] = u.UUID

	logger.Info("Renamed user")
}
<|endoftext|>