<commit_before>package netlink_test

import (
	"log"

	"github.com/mdlayher/netlink/nlenc"
	"github.com/mdlayher/netlink/nltest"

	"github.com/mdlayher/netlink"
)

// This example demonstrates using a netlink.Conn to execute requests against
// netlink.
func ExampleConn_execute() {
	// Speak to generic netlink using netlink
	const familyGeneric = 16

	c, err := netlink.Dial(familyGeneric, nil)
	if err != nil {
		log.Fatalf("failed to dial netlink: %v", err)
	}
	defer c.Close()

	// Ask netlink to send us an acknowledgement, which will contain
	// a copy of the header we sent to it
	req := netlink.Message{
		Header: netlink.Header{
			// Package netlink will automatically set header fields
			// which are set to zero
			Flags: netlink.Request | netlink.Acknowledge,
		},
	}

	// Perform a request, receive replies, and validate the replies
	msgs, err := c.Execute(req)
	if err != nil {
		log.Fatalf("failed to execute request: %v", err)
	}

	if c := len(msgs); c != 1 {
		log.Fatalf("expected 1 message, but got: %d", c)
	}

	// Decode the copied request header, starting after 4 bytes
	// indicating "success"
	var res netlink.Message
	if err := (&res).UnmarshalBinary(msgs[0].Data[4:]); err != nil {
		log.Fatalf("failed to unmarshal response: %v", err)
	}

	log.Printf("res: %+v", res)
}

// This example demonstrates using a netlink.Conn to listen for multicast group
// messages generated by the addition and deletion of network interfaces.
func ExampleConn_listenMulticast() {
	const (
		// Speak to route netlink using netlink
		familyRoute = 0

		// Listen for events triggered by addition or deletion of
		// network interfaces
		rtmGroupLink = 0x1
	)

	c, err := netlink.Dial(familyRoute, &netlink.Config{
		// Groups is a bitmask; more than one group can be specified
		// by OR'ing multiple group values together
		Groups: rtmGroupLink,
	})
	if err != nil {
		log.Fatalf("failed to dial netlink: %v", err)
	}
	defer c.Close()

	for {
		// Listen for netlink messages triggered by multicast groups
		msgs, err := c.Receive()
		if err != nil {
			log.Fatalf("failed to receive messages: %v", err)
		}

		log.Printf("msgs: %+v", msgs)
	}
}

func exampleAttributes() []byte {
	return nltest.MustMarshalAttributes([]netlink.Attribute{
		{
			Type: 1,
			Data: nlenc.Uint16Bytes(1),
		},
		{
			Type: 2,
			Data: nlenc.Bytes("hello world"),
		},
		{
			Type: 3,
			Data: nltest.MustMarshalAttributes([]netlink.Attribute{
				{
					Type: 1,
					Data: nlenc.Uint32Bytes(2),
				},
				{
					Type: 2,
					Data: nlenc.Uint32Bytes(3),
				},
			}),
		},
	})
}
<commit_msg>netlink: tidy example imports<commit_after>package netlink_test

import (
	"log"

	"github.com/mdlayher/netlink"
	"github.com/mdlayher/netlink/nlenc"
	"github.com/mdlayher/netlink/nltest"
)

// This example demonstrates using a netlink.Conn to execute requests against
// netlink.
func ExampleConn_execute() {
	// Speak to generic netlink using netlink
	const familyGeneric = 16

	c, err := netlink.Dial(familyGeneric, nil)
	if err != nil {
		log.Fatalf("failed to dial netlink: %v", err)
	}
	defer c.Close()

	// Ask netlink to send us an acknowledgement, which will contain
	// a copy of the header we sent to it
	req := netlink.Message{
		Header: netlink.Header{
			// Package netlink will automatically set header fields
			// which are set to zero
			Flags: netlink.Request | netlink.Acknowledge,
		},
	}

	// Perform a request, receive replies, and validate the replies
	msgs, err := c.Execute(req)
	if err != nil {
		log.Fatalf("failed to execute request: %v", err)
	}

	if c := len(msgs); c != 1 {
		log.Fatalf("expected 1 message, but got: %d", c)
	}

	// Decode the copied request header, starting after 4 bytes
	// indicating "success"
	var res netlink.Message
	if err := (&res).UnmarshalBinary(msgs[0].Data[4:]); err != nil {
		log.Fatalf("failed to unmarshal response: %v", err)
	}

	log.Printf("res: %+v", res)
}

// This example demonstrates using a netlink.Conn to listen for multicast group
// messages generated by the addition and deletion of network interfaces.
func ExampleConn_listenMulticast() {
	const (
		// Speak to route netlink using netlink
		familyRoute = 0

		// Listen for events triggered by addition or deletion of
		// network interfaces
		rtmGroupLink = 0x1
	)

	c, err := netlink.Dial(familyRoute, &netlink.Config{
		// Groups is a bitmask; more than one group can be specified
		// by OR'ing multiple group values together
		Groups: rtmGroupLink,
	})
	if err != nil {
		log.Fatalf("failed to dial netlink: %v", err)
	}
	defer c.Close()

	for {
		// Listen for netlink messages triggered by multicast groups
		msgs, err := c.Receive()
		if err != nil {
			log.Fatalf("failed to receive messages: %v", err)
		}

		log.Printf("msgs: %+v", msgs)
	}
}

func exampleAttributes() []byte {
	return nltest.MustMarshalAttributes([]netlink.Attribute{
		{
			Type: 1,
			Data: nlenc.Uint16Bytes(1),
		},
		{
			Type: 2,
			Data: nlenc.Bytes("hello world"),
		},
		{
			Type: 3,
			Data: nltest.MustMarshalAttributes([]netlink.Attribute{
				{
					Type: 1,
					Data: nlenc.Uint32Bytes(2),
				},
				{
					Type: 2,
					Data: nlenc.Uint32Bytes(3),
				},
			}),
		},
	})
}
<|endoftext|>


<commit_before>// +build !windows

package godirwalk

import (
	"os"
	"syscall"
	"unsafe"
)

// MinimumScratchBufferSize specifies the minimum size of the scratch buffer
// that Walk, ReadDirents, ReadDirnames, and Scandir will use when reading file
// entries from the operating system. It is initialized to the result from
// calling `os.Getpagesize()` during program startup.
var MinimumScratchBufferSize = os.Getpagesize()

// Scanner is an iterator to enumerate the contents of a directory.
type Scanner struct {
	scratchBuffer []byte // read directory bytes from file system into this buffer
	workBuffer    []byte // points into scratchBuffer, from which we chunk out directory entries
	osDirname     string
	childName     string
	err           error    // err is the error associated with scanning directory
	statErr       error    // statErr is any error return while attempting to stat an entry
	dh            *os.File // used to close directory after done reading
	de            *Dirent  // most recently decoded directory entry
	sde           *syscall.Dirent
	fd            int // file descriptor used to read entries from directory
}

// NewScanner returns a new directory Scanner.
func NewScanner(osDirname string) (*Scanner, error) {
	dh, err := os.Open(osDirname)
	if err != nil {
		return nil, err
	}
	scanner := &Scanner{
		scratchBuffer: make([]byte, MinimumScratchBufferSize),
		osDirname:     osDirname,
		dh:            dh,
		fd:            int(dh.Fd()),
	}
	return scanner, nil
}

// Dirent returns the current directory entry while scanning a directory.
func (s *Scanner) Dirent() (*Dirent, error) {
	if s.de == nil {
		s.de = &Dirent{name: s.childName}
		s.de.modeType, s.statErr = modeTypeFromDirent(s.sde, s.osDirname, s.childName)
	}
	return s.de, s.statErr
}

// done is called when directory scanner unable to continue, with either the
// triggering error, or nil when there are simply no more entries to read from
// the directory.
func (s *Scanner) done(err error) {
	if s.dh == nil {
		return
	}
	cerr := s.dh.Close()
	s.dh = nil

	if err == nil {
		s.err = cerr
	} else {
		s.err = err
	}

	s.osDirname, s.childName = "", ""
	s.scratchBuffer, s.workBuffer = nil, nil
	s.statErr, s.de, s.sde = nil, nil, nil
	s.fd = 0
}

// Err returns the error associated with scanning a directory.
func (s *Scanner) Err() error {
	s.done(s.err)
	return s.err
}

// Name returns the name of the current directory entry while scanning a
// directory.
func (s *Scanner) Name() string { return s.childName }

// Scan potentially reads and then decodes the next directory entry from the
// file system.
//
// When it returns false, this releases resources used by the Scanner then
// returns any error associated with closing the file system directory resource.
func (s *Scanner) Scan() bool {
	if s.err != nil {
		return false
	}

	for {
		// When the work buffer has nothing remaining to decode, we need to load
		// more data from disk.
		if len(s.workBuffer) == 0 {
			n, err := syscall.ReadDirent(s.fd, s.scratchBuffer)
			if err != nil {
				s.done(err)
				return false
			}
			if n <= 0 { // end of directory
				s.done(nil)
				return false
			}
			s.workBuffer = s.scratchBuffer[:n] // trim work buffer to number of bytes read
		}

		// Loop until we have a usable file system entry, or we run out of data
		// in the work buffer.
		for len(s.workBuffer) > 0 {
			s.sde = (*syscall.Dirent)(unsafe.Pointer(&s.workBuffer[0])) // point entry to first syscall.Dirent in buffer
			s.workBuffer = s.workBuffer[s.sde.Reclen:]                  // advance buffer for next iteration through loop

			if inoFromDirent(s.sde) == 0 {
				continue // inode set to 0 indicates an entry that was marked as deleted
			}

			nameSlice := nameFromDirent(s.sde)
			namlen := len(nameSlice)
			if (namlen == 0) || (namlen == 1 && nameSlice[0] == '.') || (namlen == 2 && nameSlice[0] == '.' && nameSlice[1] == '.') {
				continue // skip unimportant entries
			}

			s.de = nil
			s.childName = string(nameSlice)
			return true
		}
		// No more data in the work buffer, so loop around in the outside loop
		// to fetch more data.
	}
}
<commit_msg>Scanner on UNIX internally stores io.EOF when done<commit_after>// +build !windows

package godirwalk

import (
	"io"
	"os"
	"syscall"
	"unsafe"
)

// MinimumScratchBufferSize specifies the minimum size of the scratch buffer
// that Walk, ReadDirents, ReadDirnames, and Scandir will use when reading file
// entries from the operating system. It is initialized to the result from
// calling `os.Getpagesize()` during program startup.
var MinimumScratchBufferSize = os.Getpagesize()

// Scanner is an iterator to enumerate the contents of a directory.
type Scanner struct {
	scratchBuffer []byte // read directory bytes from file system into this buffer
	workBuffer    []byte // points into scratchBuffer, from which we chunk out directory entries
	osDirname     string
	childName     string
	err           error    // err is the error associated with scanning directory
	statErr       error    // statErr is any error return while attempting to stat an entry
	dh            *os.File // used to close directory after done reading
	de            *Dirent  // most recently decoded directory entry
	sde           *syscall.Dirent
	fd            int // file descriptor used to read entries from directory
}

// NewScanner returns a new directory Scanner.
func NewScanner(osDirname string) (*Scanner, error) {
	dh, err := os.Open(osDirname)
	if err != nil {
		return nil, err
	}
	scanner := &Scanner{
		scratchBuffer: make([]byte, MinimumScratchBufferSize),
		osDirname:     osDirname,
		dh:            dh,
		fd:            int(dh.Fd()),
	}
	return scanner, nil
}

// Dirent returns the current directory entry while scanning a directory.
func (s *Scanner) Dirent() (*Dirent, error) {
	if s.de == nil {
		s.de = &Dirent{name: s.childName}
		s.de.modeType, s.statErr = modeTypeFromDirent(s.sde, s.osDirname, s.childName)
	}
	return s.de, s.statErr
}

// done is called when directory scanner unable to continue, with either the
// triggering error, or nil when there are simply no more entries to read from
// the directory.
func (s *Scanner) done(err error) {
	if s.dh == nil {
		return
	}
	cerr := s.dh.Close()
	s.dh = nil

	if err == nil {
		s.err = cerr
	} else {
		s.err = err
	}

	s.osDirname, s.childName = "", ""
	s.scratchBuffer, s.workBuffer = nil, nil
	s.statErr, s.de, s.sde = nil, nil, nil
	s.fd = 0
}

// Err returns the error associated with scanning a directory.
func (s *Scanner) Err() error {
	s.done(s.err)
	if s.err == io.EOF {
		return nil
	}
	return s.err
}

// Name returns the name of the current directory entry while scanning a
// directory.
func (s *Scanner) Name() string { return s.childName }

// Scan potentially reads and then decodes the next directory entry from the
// file system.
//
// When it returns false, this releases resources used by the Scanner then
// returns any error associated with closing the file system directory resource.
func (s *Scanner) Scan() bool {
	if s.err != nil {
		return false
	}

	for {
		// When the work buffer has nothing remaining to decode, we need to load
		// more data from disk.
		if len(s.workBuffer) == 0 {
			n, err := syscall.ReadDirent(s.fd, s.scratchBuffer)
			if err != nil {
				s.done(err)
				return false
			}
			if n <= 0 { // end of directory
				s.done(io.EOF)
				return false
			}
			s.workBuffer = s.scratchBuffer[:n] // trim work buffer to number of bytes read
		}

		// Loop until we have a usable file system entry, or we run out of data
		// in the work buffer.
		for len(s.workBuffer) > 0 {
			s.sde = (*syscall.Dirent)(unsafe.Pointer(&s.workBuffer[0])) // point entry to first syscall.Dirent in buffer
			s.workBuffer = s.workBuffer[s.sde.Reclen:]                  // advance buffer for next iteration through loop

			if inoFromDirent(s.sde) == 0 {
				continue // inode set to 0 indicates an entry that was marked as deleted
			}

			nameSlice := nameFromDirent(s.sde)
			namlen := len(nameSlice)
			if (namlen == 0) || (namlen == 1 && nameSlice[0] == '.') || (namlen == 2 && nameSlice[0] == '.' && nameSlice[1] == '.') {
				continue // skip unimportant entries
			}

			s.de = nil
			s.childName = string(nameSlice)
			return true
		}
		// No more data in the work buffer, so loop around in the outside loop
		// to fetch more data.
	}
}
<|endoftext|>


<commit_before>package scanner

import (
	"bytes"
	"io/ioutil"
	"log"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"

	"github.com/calmh/syncthing/protocol"
)

type Walker struct {
	// Dir is the base directory for the walk
	Dir string
	// If FollowSymlinks is true, symbolic links directly under Dir will be followed.
	// Symbolic links at deeper levels are never followed regardless of this flag.
	FollowSymlinks bool
	// BlockSize controls the size of the block used when hashing.
	BlockSize int
	// If IgnoreFile is not empty, it is the name used for the file that holds ignore patterns.
	IgnoreFile string
	// If TempNamer is not nil, it is used to ignore tempory files when walking.
	TempNamer TempNamer
	// If Suppressor is not nil, it is queried for supression of modified files.
	Suppressor Suppressor

	previous   map[string]File // file name -> last seen file state
	suppressed map[string]bool // file name -> suppression status
}

type TempNamer interface {
	// Temporary returns a temporary name for the filed referred to by path.
	TempName(path string) string
	// IsTemporary returns true if path refers to the name of temporary file.
	IsTemporary(path string) bool
}

type Suppressor interface {
	// Supress returns true if the update to the named file should be ignored.
	Suppress(name string, fi os.FileInfo) bool
}

// Walk returns the list of files found in the local repository by scanning the
// file system. Files are blockwise hashed.
func (w *Walker) Walk() (files []File, ignore map[string][]string) {
	w.lazyInit()

	if debug {
		dlog.Println("Walk", w.Dir, w.FollowSymlinks, w.BlockSize, w.IgnoreFile)
	}
	t0 := time.Now()

	ignore = make(map[string][]string)
	hashFiles := w.walkAndHashFiles(&files, ignore)

	filepath.Walk(w.Dir, w.loadIgnoreFiles(w.Dir, ignore))
	filepath.Walk(w.Dir, hashFiles)

	if w.FollowSymlinks {
		d, err := os.Open(w.Dir)
		if err != nil {
			return
		}
		defer d.Close()

		fis, err := d.Readdir(-1)
		if err != nil {
			return
		}

		for _, info := range fis {
			if info.Mode()&os.ModeSymlink != 0 {
				dir := path.Join(w.Dir, info.Name()) + "/"
				filepath.Walk(dir, w.loadIgnoreFiles(dir, ignore))
				filepath.Walk(dir, hashFiles)
			}
		}
	}

	if debug {
		t1 := time.Now()
		d := t1.Sub(t0).Seconds()
		dlog.Printf("Walk in %.02f ms, %.0f files/s", d*1000, float64(len(files))/d)
	}
	return
}

// CleanTempFiles removes all files that match the temporary filename pattern.
func (w *Walker) CleanTempFiles() {
	filepath.Walk(w.Dir, w.cleanTempFile)
}

func (w *Walker) lazyInit() {
	if w.previous == nil {
		w.previous = make(map[string]File)
		w.suppressed = make(map[string]bool)
	}
}

func (w *Walker) loadIgnoreFiles(dir string, ign map[string][]string) filepath.WalkFunc {
	return func(p string, info os.FileInfo, err error) error {
		if err != nil {
			return nil
		}

		p, err = filepath.Rel(dir, p)
		if err != nil {
			return nil
		}

		if pn, sn := path.Split(p); sn == w.IgnoreFile {
			pn := strings.Trim(pn, "/")
			bs, _ := ioutil.ReadFile(p)
			lines := bytes.Split(bs, []byte("\n"))
			var patterns []string
			for _, line := range lines {
				if len(line) > 0 {
					patterns = append(patterns, string(line))
				}
			}
			ign[pn] = patterns
		}

		return nil
	}
}

func (w *Walker) walkAndHashFiles(res *[]File, ign map[string][]string) filepath.WalkFunc {
	return func(p string, info os.FileInfo, err error) error {
		if err != nil {
			if debug {
				dlog.Println("error:", p, info, err)
			}
			return nil
		}

		p, err = filepath.Rel(w.Dir, p)
		if err != nil {
			return nil
		}

		if w.TempNamer != nil && w.TempNamer.IsTemporary(p) {
			if debug {
				dlog.Println("temporary:", p)
			}
			return nil
		}

		if _, sn := path.Split(p); sn == w.IgnoreFile {
			if debug {
				dlog.Println("ignorefile:", p)
			}
			return nil
		}

		if w.ignoreFile(ign, p) {
			if debug {
				dlog.Println("ignored:", p)
			}
			return nil
		}

		if info.Mode()&os.ModeType == 0 {
			modified := info.ModTime().Unix()
			pf := w.previous[p]

			if pf.Modified == modified {
				if nf := uint32(info.Mode()); nf != pf.Flags {
					if debug {
						dlog.Println("new flags:", p)
					}
					pf.Flags = nf
					pf.Version++
					w.previous[p] = pf
				} else if debug {
					dlog.Println("unchanged:", p)
				}
				*res = append(*res, pf)
				return nil
			}

			if w.Suppressor != nil && w.Suppressor.Suppress(p, info) {
				if debug {
					dlog.Println("suppressed:", p)
				}
				if !w.suppressed[p] {
					w.suppressed[p] = true
					log.Printf("INFO: Changes to %q are being temporarily suppressed because it changes too frequently.", p)
				}
				f := pf
				f.Flags = protocol.FlagInvalid
				f.Blocks = nil
				*res = append(*res, f)
				return nil
			} else if w.suppressed[p] {
				log.Printf("INFO: Changes to %q are no longer suppressed.", p)
				delete(w.suppressed, p)
			}

			fd, err := os.Open(p)
			if err != nil {
				return nil
			}
			defer fd.Close()

			t0 := time.Now()
			blocks, err := Blocks(fd, w.BlockSize)
			if err != nil {
				if debug {
					dlog.Println("hash error:", p, err)
				}
				return nil
			}
			if debug {
				t1 := time.Now()
				dlog.Println("hashed:", p, ";", len(blocks), "blocks;", info.Size(), "bytes;", int(float64(info.Size())/1024/t1.Sub(t0).Seconds()), "KB/s")
			}
			f := File{
				Name:     p,
				Size:     info.Size(),
				Flags:    uint32(info.Mode()),
				Modified: modified,
				Blocks:   blocks,
			}
			w.previous[p] = f
			*res = append(*res, f)
		}

		return nil
	}
}

func (w *Walker) cleanTempFile(path string, info os.FileInfo, err error) error {
	if err != nil {
		return err
	}
	if info.Mode()&os.ModeType == 0 && w.TempNamer.IsTemporary(path) {
		os.Remove(path)
	}
	return nil
}

func (w *Walker) ignoreFile(patterns map[string][]string, file string) bool {
	first, last := path.Split(file)
	for prefix, pats := range patterns {
		if len(prefix) == 0 || prefix == first || strings.HasPrefix(first, prefix+"/") {
			for _, pattern := range pats {
				if match, _ := path.Match(pattern, last); match {
					return true
				}
			}
		}
	}
	return false
}
<commit_msg>Fix relative path open bug<commit_after>package scanner

import (
	"bytes"
	"io/ioutil"
	"log"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"

	"github.com/calmh/syncthing/protocol"
)

type Walker struct {
	// Dir is the base directory for the walk
	Dir string
	// If FollowSymlinks is true, symbolic links directly under Dir will be followed.
	// Symbolic links at deeper levels are never followed regardless of this flag.
	FollowSymlinks bool
	// BlockSize controls the size of the block used when hashing.
	BlockSize int
	// If IgnoreFile is not empty, it is the name used for the file that holds ignore patterns.
	IgnoreFile string
	// If TempNamer is not nil, it is used to ignore tempory files when walking.
	TempNamer TempNamer
	// If Suppressor is not nil, it is queried for supression of modified files.
	Suppressor Suppressor

	previous   map[string]File // file name -> last seen file state
	suppressed map[string]bool // file name -> suppression status
}

type TempNamer interface {
	// Temporary returns a temporary name for the filed referred to by path.
	TempName(path string) string
	// IsTemporary returns true if path refers to the name of temporary file.
	IsTemporary(path string) bool
}

type Suppressor interface {
	// Supress returns true if the update to the named file should be ignored.
	Suppress(name string, fi os.FileInfo) bool
}

// Walk returns the list of files found in the local repository by scanning the
// file system. Files are blockwise hashed.
func (w *Walker) Walk() (files []File, ignore map[string][]string) {
	w.lazyInit()

	if debug {
		dlog.Println("Walk", w.Dir, w.FollowSymlinks, w.BlockSize, w.IgnoreFile)
	}
	t0 := time.Now()

	ignore = make(map[string][]string)
	hashFiles := w.walkAndHashFiles(&files, ignore)

	filepath.Walk(w.Dir, w.loadIgnoreFiles(w.Dir, ignore))
	filepath.Walk(w.Dir, hashFiles)

	if w.FollowSymlinks {
		d, err := os.Open(w.Dir)
		if err != nil {
			return
		}
		defer d.Close()

		fis, err := d.Readdir(-1)
		if err != nil {
			return
		}

		for _, info := range fis {
			if info.Mode()&os.ModeSymlink != 0 {
				dir := path.Join(w.Dir, info.Name()) + "/"
				filepath.Walk(dir, w.loadIgnoreFiles(dir, ignore))
				filepath.Walk(dir, hashFiles)
			}
		}
	}

	if debug {
		t1 := time.Now()
		d := t1.Sub(t0).Seconds()
		dlog.Printf("Walk in %.02f ms, %.0f files/s", d*1000, float64(len(files))/d)
	}
	return
}

// CleanTempFiles removes all files that match the temporary filename pattern.
func (w *Walker) CleanTempFiles() {
	filepath.Walk(w.Dir, w.cleanTempFile)
}

func (w *Walker) lazyInit() {
	if w.previous == nil {
		w.previous = make(map[string]File)
		w.suppressed = make(map[string]bool)
	}
}

func (w *Walker) loadIgnoreFiles(dir string, ign map[string][]string) filepath.WalkFunc {
	return func(p string, info os.FileInfo, err error) error {
		if err != nil {
			return nil
		}

		rn, err := filepath.Rel(dir, p)
		if err != nil {
			return nil
		}

		if pn, sn := path.Split(rn); sn == w.IgnoreFile {
			pn := strings.Trim(pn, "/")
			bs, _ := ioutil.ReadFile(p)
			lines := bytes.Split(bs, []byte("\n"))
			var patterns []string
			for _, line := range lines {
				if len(line) > 0 {
					patterns = append(patterns, string(line))
				}
			}
			ign[pn] = patterns
		}

		return nil
	}
}

func (w *Walker) walkAndHashFiles(res *[]File, ign map[string][]string) filepath.WalkFunc {
	return func(p string, info os.FileInfo, err error) error {
		if err != nil {
			if debug {
				dlog.Println("error:", p, info, err)
			}
			return nil
		}

		rn, err := filepath.Rel(w.Dir, p)
		if err != nil {
			if debug {
				dlog.Println("rel error:", p, err)
			}
			return nil
		}

		if w.TempNamer != nil && w.TempNamer.IsTemporary(rn) {
			if debug {
				dlog.Println("temporary:", rn)
			}
			return nil
		}

		if _, sn := path.Split(rn); sn == w.IgnoreFile {
			if debug {
				dlog.Println("ignorefile:", rn)
			}
			return nil
		}

		if w.ignoreFile(ign, rn) {
			if debug {
				dlog.Println("ignored:", rn)
			}
			return nil
		}

		if info.Mode()&os.ModeType == 0 {
			modified := info.ModTime().Unix()
			pf := w.previous[rn]

			if pf.Modified == modified {
				if nf := uint32(info.Mode()); nf != pf.Flags {
					if debug {
						dlog.Println("new flags:", rn)
					}
					pf.Flags = nf
					pf.Version++
					w.previous[rn] = pf
				} else if debug {
					dlog.Println("unchanged:", p)
				}
				*res = append(*res, pf)
				return nil
			}

			if w.Suppressor != nil && w.Suppressor.Suppress(rn, info) {
				if debug {
					dlog.Println("suppressed:", rn)
				}
				if !w.suppressed[rn] {
					w.suppressed[rn] = true
					log.Printf("INFO: Changes to %q are being temporarily suppressed because it changes too frequently.", p)
				}
				f := pf
				f.Flags = protocol.FlagInvalid
				f.Blocks = nil
				*res = append(*res, f)
				return nil
			} else if w.suppressed[rn] {
				log.Printf("INFO: Changes to %q are no longer suppressed.", p)
				delete(w.suppressed, rn)
			}

			fd, err := os.Open(p)
			if err != nil {
				if debug {
					dlog.Println("open:", p, err)
				}
				return nil
			}
			defer fd.Close()

			t0 := time.Now()
			blocks, err := Blocks(fd, w.BlockSize)
			if err != nil {
				if debug {
					dlog.Println("hash error:", rn, err)
				}
				return nil
			}
			if debug {
				t1 := time.Now()
				dlog.Println("hashed:", rn, ";", len(blocks), "blocks;", info.Size(), "bytes;", int(float64(info.Size())/1024/t1.Sub(t0).Seconds()), "KB/s")
			}
			f := File{
				Name:     rn,
				Size:     info.Size(),
				Flags:    uint32(info.Mode()),
				Modified: modified,
				Blocks:   blocks,
			}
			w.previous[rn] = f
			*res = append(*res, f)
		}

		return nil
	}
}

func (w *Walker) cleanTempFile(path string, info os.FileInfo, err error) error {
	if err != nil {
		return err
	}
	if info.Mode()&os.ModeType == 0 && w.TempNamer.IsTemporary(path) {
		os.Remove(path)
	}
	return nil
}

func (w *Walker) ignoreFile(patterns map[string][]string, file string) bool {
	first, last := path.Split(file)
	for prefix, pats := range patterns {
		if len(prefix) == 0 || prefix == first || strings.HasPrefix(first, prefix+"/") {
			for _, pattern := range pats {
				if match, _ := path.Match(pattern, last); match {
					return true
				}
			}
		}
	}
	return false
}
<|endoftext|>


<commit_before>package tar

import (
	"os"
	"path/filepath"
	"testing"

	"github.com/polydawn/gosh"
	. "github.com/smartystreets/goconvey/convey"
	"polydawn.net/repeatr/def"
	"polydawn.net/repeatr/lib/fspatch"
	"polydawn.net/repeatr/lib/testutil"
	"polydawn.net/repeatr/lib/testutil/filefixture"
	"polydawn.net/repeatr/rio"
)

const ubuntuTarballHash = "uJRF46th6rYHt0zt_n3fcDuBfGFVPS6lzRZla5hv6iDoh5DVVzxUTMMzENfPoboL"

func TestTarInputCompat(t *testing.T) {
	projPath, _ := os.Getwd()
	projPath = filepath.Dir(filepath.Dir(filepath.Dir(projPath)))

	Convey("Unpacking tars should match exec untar", t,
		testutil.Requires(testutil.RequiresRoot, testutil.WithTmpdir(func(c C) {
			checkEquivalence := func(hash, filename string, paveBase bool) {
				transmat := New("./workdir/tar")

				// apply it; hope it doesn't blow up
				arena := transmat.Materialize(
					rio.TransmatKind("tar"),
					rio.CommitID(hash),
					[]rio.SiloURI{
						rio.SiloURI("file://" + filename),
					},
					testutil.TestLogger(c),
				)
				defer arena.Teardown()

				// do a native untar; since we don't have an upfront fixture
				//  for this thing, we'll compare the two as filesystems.
				// this is not well isolated from the host; consider improving that a todo.
				os.Mkdir("./untar", 0755)
				tarProc := gosh.Gosh(
					"tar",
					"-xf", filename,
					"-C", "./untar",
					gosh.NullIO,
				).RunAndReport()
				So(tarProc.GetExitCode(), ShouldEqual, 0)
				// native untar may or may not have an opinion about the base dir, depending on how it was formed.
				// but our scans do, so, if the `paveBase` flag was set to warn us that the tar was missing an "./" entry, flatten that here.
				if paveBase {
					So(fspatch.LUtimesNano("./untar", def.Epochwhen, def.Epochwhen), ShouldBeNil)
				}

				// scan and compare
				scan1 := filefixture.Scan(arena.Path())
				scan2 := filefixture.Scan("./untar")
				// boy, that's entertaining though: gnu tar does all the same stuff,
				//  except it doesn't honor our nanosecond timings.
				// also exclude bodies because they're *big*.
				comparisonLevel := filefixture.CompareDefaults &^ filefixture.CompareSubsecond &^ filefixture.CompareBody
				So(scan1.Describe(comparisonLevel), ShouldEqual, scan2.Describe(comparisonLevel))
			}

			Convey("Given a fixture tarball complete with base dir", func() {
				checkEquivalence(
					"BX0jm4jRNCg1KMbZfv4zp7ZaShx9SUXKaDrO-Xy6mWIoWOCFP5VnDHDDR3nU4PrR",
					filepath.Join(projPath, "data/fixture/tar_withBase.tgz"),
					false,
				)
			})

			Convey("Given a fixture tarball lacking base dir", func() {
				checkEquivalence(
					"ZdV3xhCGWeJmsfeHpDF4nF9stwvdskYwcepKMcOf7a2ziax1YGjQvGTJjRWFkvG1",
					filepath.Join(projPath, "data/fixture/tar_sansBase.tgz"),
					true,
				)
			})

			Convey("Given a fixture tarball containing ubuntu",
				testutil.Requires(testutil.RequiresLongRun, func() {
					checkEquivalence(
						ubuntuTarballHash,
						filepath.Join(projPath, "assets/ubuntu.tar.gz"),
						true,
					)
				}),
			)
		})),
	)

	Convey("Bouncing unusual tars should match hash", t,
		// where really all "unusual" means is "valid tar, but not from our own cannonical output".
		testutil.Requires(
			testutil.RequiresRoot,
			testutil.WithTmpdir(func(c C) {
				checkBounce := func(hash, filename string) {
					transmat := New("./workdir/tar")

					// apply it; hope it doesn't blow up
					arena := transmat.Materialize(
						rio.TransmatKind("tar"),
						rio.CommitID(hash),
						[]rio.SiloURI{
							rio.SiloURI("file://" + filename),
						},
						testutil.TestLogger(c),
					)
					defer arena.Teardown()

					// scan and compare
					commitID := transmat.Scan(rio.TransmatKind("tar"), arena.Path(), nil, testutil.TestLogger(c))
					// debug: gosh.Sh("tar", "--utc", "-xOvf", filename)
					So(commitID, ShouldEqual, rio.CommitID(hash))

				}

				Convey("Given a fixture tarball complete with base dir", func() {
					checkBounce(
						"BX0jm4jRNCg1KMbZfv4zp7ZaShx9SUXKaDrO-Xy6mWIoWOCFP5VnDHDDR3nU4PrR",
						filepath.Join(projPath, "data/fixture/tar_withBase.tgz"),
					)
				})

				Convey("Given a fixture tarball lacking base dir", func() {
					checkBounce(
						"ZdV3xhCGWeJmsfeHpDF4nF9stwvdskYwcepKMcOf7a2ziax1YGjQvGTJjRWFkvG1",
						filepath.Join(projPath, "data/fixture/tar_sansBase.tgz"),
					)
				})

				// this won't fly until we support hardlinks; the original asset uses them.
				//	Convey("Given a fixture tarball containing ubuntu",
				//		testutil.Requires(testutil.RequiresLongRun, func() {
				//			checkBounce(
				//				ubuntuTarballHash,
				//				filepath.Join(projPath, "assets/ubuntu.tar.gz"),
				//			)
				//		}),
				//	)
			}),
		),
	)
}
<commit_msg>Remove reference to old(oldold) assets from tests.<commit_after>package tar

import (
	"os"
	"path/filepath"
	"testing"

	"github.com/polydawn/gosh"
	. "github.com/smartystreets/goconvey/convey"
	"polydawn.net/repeatr/def"
	"polydawn.net/repeatr/lib/fspatch"
	"polydawn.net/repeatr/lib/testutil"
	"polydawn.net/repeatr/lib/testutil/filefixture"
	"polydawn.net/repeatr/rio"
)

func TestTarInputCompat(t *testing.T) {
	projPath, _ := os.Getwd()
	projPath = filepath.Dir(filepath.Dir(filepath.Dir(projPath)))

	Convey("Unpacking tars should match exec untar", t,
		testutil.Requires(testutil.RequiresRoot, testutil.WithTmpdir(func(c C) {
			checkEquivalence := func(hash, filename string, paveBase bool) {
				transmat := New("./workdir/tar")

				// apply it; hope it doesn't blow up
				arena := transmat.Materialize(
					rio.TransmatKind("tar"),
					rio.CommitID(hash),
					[]rio.SiloURI{
						rio.SiloURI("file://" + filename),
					},
					testutil.TestLogger(c),
				)
				defer arena.Teardown()

				// do a native untar; since we don't have an upfront fixture
				//  for this thing, we'll compare the two as filesystems.
				// this is not well isolated from the host; consider improving that a todo.
				os.Mkdir("./untar", 0755)
				tarProc := gosh.Gosh(
					"tar",
					"-xf", filename,
					"-C", "./untar",
					gosh.NullIO,
				).RunAndReport()
				So(tarProc.GetExitCode(), ShouldEqual, 0)
				// native untar may or may not have an opinion about the base dir, depending on how it was formed.
				// but our scans do, so, if the `paveBase` flag was set to warn us that the tar was missing an "./" entry, flatten that here.
				if paveBase {
					So(fspatch.LUtimesNano("./untar", def.Epochwhen, def.Epochwhen), ShouldBeNil)
				}

				// scan and compare
				scan1 := filefixture.Scan(arena.Path())
				scan2 := filefixture.Scan("./untar")
				// boy, that's entertaining though: gnu tar does all the same stuff,
				//  except it doesn't honor our nanosecond timings.
				// also exclude bodies because they're *big*.
				comparisonLevel := filefixture.CompareDefaults &^ filefixture.CompareSubsecond &^ filefixture.CompareBody
				So(scan1.Describe(comparisonLevel), ShouldEqual, scan2.Describe(comparisonLevel))
			}

			Convey("Given a fixture tarball complete with base dir", func() {
				checkEquivalence(
					"BX0jm4jRNCg1KMbZfv4zp7ZaShx9SUXKaDrO-Xy6mWIoWOCFP5VnDHDDR3nU4PrR",
					filepath.Join(projPath, "data/fixture/tar_withBase.tgz"),
					false,
				)
			})

			Convey("Given a fixture tarball lacking base dir", func() {
				checkEquivalence(
					"ZdV3xhCGWeJmsfeHpDF4nF9stwvdskYwcepKMcOf7a2ziax1YGjQvGTJjRWFkvG1",
					filepath.Join(projPath, "data/fixture/tar_sansBase.tgz"),
					true,
				)
			})
		})),
	)

	Convey("Bouncing unusual tars should match hash", t,
		// where really all "unusual" means is "valid tar, but not from our own cannonical output".
		testutil.Requires(
			testutil.RequiresRoot,
			testutil.WithTmpdir(func(c C) {
				checkBounce := func(hash, filename string) {
					transmat := New("./workdir/tar")

					// apply it; hope it doesn't blow up
					arena := transmat.Materialize(
						rio.TransmatKind("tar"),
						rio.CommitID(hash),
						[]rio.SiloURI{
							rio.SiloURI("file://" + filename),
						},
						testutil.TestLogger(c),
					)
					defer arena.Teardown()

					// scan and compare
					commitID := transmat.Scan(rio.TransmatKind("tar"), arena.Path(), nil, testutil.TestLogger(c))
					// debug: gosh.Sh("tar", "--utc", "-xOvf", filename)
					So(commitID, ShouldEqual, rio.CommitID(hash))

				}

				Convey("Given a fixture tarball complete with base dir", func() {
					checkBounce(
						"BX0jm4jRNCg1KMbZfv4zp7ZaShx9SUXKaDrO-Xy6mWIoWOCFP5VnDHDDR3nU4PrR",
						filepath.Join(projPath, "data/fixture/tar_withBase.tgz"),
					)
				})

				Convey("Given a fixture tarball lacking base dir", func() {
					checkBounce(
						"ZdV3xhCGWeJmsfeHpDF4nF9stwvdskYwcepKMcOf7a2ziax1YGjQvGTJjRWFkvG1",
						filepath.Join(projPath, "data/fixture/tar_sansBase.tgz"),
					)
				})
			}),
		),
	)
}
<|endoftext|>


<commit_before>package google

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"runtime"
	"strings"

	// TODO(dcunnin): Use version code from version.go
	// "github.com/hashicorp/terraform"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	"golang.org/x/oauth2/jwt"
	"google.golang.org/api/compute/v1"
	"google.golang.org/api/container/v1"
	"google.golang.org/api/dns/v1"
	"google.golang.org/api/storage/v1"
)

// Config is the configuration structure used to instantiate the Google
// provider.
type Config struct {
	AccountFile string
	Project     string
	Region      string

	clientCompute *compute.Service
	clientContainer *container.Service
	clientDns *dns.Service
	clientStorage *storage.Service
}

func (c *Config) loadAndValidate() error {
	var account accountFile

	if c.AccountFile == "" {
		c.AccountFile = os.Getenv("GOOGLE_ACCOUNT_FILE")
	}
	if c.Project == "" {
		c.Project = os.Getenv("GOOGLE_PROJECT")
	}
	if c.Region == "" {
		c.Region = os.Getenv("GOOGLE_REGION")
	}

	var client *http.Client

	if c.AccountFile != "" {
		contents := c.AccountFile

		// Assume account_file is a JSON string
		if err := parseJSON(&account, contents); err != nil {
			// If account_file was not JSON, assume it is a file path instead
			if _, err := os.Stat(c.AccountFile); os.IsNotExist(err) {
				return fmt.Errorf(
					"account_file path does not exist: %s",
					c.AccountFile)
			}

			b, err := ioutil.ReadFile(c.AccountFile)
			if err != nil {
				return fmt.Errorf(
					"Error reading account_file from path '%s': %s",
					c.AccountFile,
					err)
			}

			contents = string(b)
		}

		if err := parseJSON(&account, contents); err != nil {
			return fmt.Errorf(
				"Error parsing account file '%s': %s",
				contents,
				err)
		}

		clientScopes := []string{
			"https://www.googleapis.com/auth/compute",
			"https://www.googleapis.com/auth/cloud-platform",
			"https://www.googleapis.com/auth/ndev.clouddns.readwrite",
			"https://www.googleapis.com/auth/devstorage.full_control",
		}

		// Get the token for use in our requests
		log.Printf("[INFO] Requesting Google token...")
		log.Printf("[INFO]   -- Email: %s", account.ClientEmail)
		log.Printf("[INFO]   -- Scopes: %s", clientScopes)
		log.Printf("[INFO]   -- Private Key Length: %d", len(account.PrivateKey))

		conf := jwt.Config{
			Email:      account.ClientEmail,
			PrivateKey: []byte(account.PrivateKey),
			Scopes:     clientScopes,
			TokenURL:   "https://accounts.google.com/o/oauth2/token",
		}

		// Initiate an http.Client. The following GET request will be
		// authorized and authenticated on the behalf of
		// your service account.
		client = conf.Client(oauth2.NoContext)

	} else {
		log.Printf("[INFO] Requesting Google token via GCE Service Role...")
		client = &http.Client{
			Transport: &oauth2.Transport{
				// Fetch from Google Compute Engine's metadata server to retrieve
				// an access token for the provided account.
				// If no account is specified, "default" is used.
				Source: google.ComputeTokenSource(""),
			},
		}

	}

	// Build UserAgent
	versionString := "0.0.0"
	// TODO(dcunnin): Use Terraform's version code from version.go
	// versionString := main.Version
	// if main.VersionPrerelease != "" {
	// 	versionString = fmt.Sprintf("%s-%s", versionString, main.VersionPrerelease)
	// }
	userAgent := fmt.Sprintf(
		"(%s %s) Terraform/%s", runtime.GOOS, runtime.GOARCH, versionString)

	var err error

	log.Printf("[INFO] Instantiating GCE client...")
	c.clientCompute, err = compute.New(client)
	if err != nil {
		return err
	}
	c.clientCompute.UserAgent = userAgent

	log.Printf("[INFO] Instantiating GKE client...")
	c.clientContainer, err = container.New(client)
	if err != nil {
		return err
	}
	c.clientContainer.UserAgent = userAgent

	log.Printf("[INFO] Instantiating Google Cloud DNS client...")
	c.clientDns, err = dns.New(client)
	if err != nil {
		return err
	}
	c.clientDns.UserAgent = userAgent

	log.Printf("[INFO] Instantiating Google Storage Client...")
	c.clientStorage, err = storage.New(client)
	if err != nil {
		return err
	}
	c.clientStorage.UserAgent = userAgent

	return nil
}

// accountFile represents the structure of the account file JSON file.
type accountFile struct {
	PrivateKeyId string `json:"private_key_id"`
	PrivateKey   string `json:"private_key"`
	ClientEmail  string `json:"client_email"`
	ClientId     string `json:"client_id"`
}

func parseJSON(result interface{}, contents string) error {
	r := strings.NewReader(contents)
	dec := json.NewDecoder(r)

	return dec.Decode(result)
}
<commit_msg>providers/google: Fix reading account_file path<commit_after>package google

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"runtime"
	"strings"

	// TODO(dcunnin): Use version code from version.go
	// "github.com/hashicorp/terraform"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	"golang.org/x/oauth2/jwt"
	"google.golang.org/api/compute/v1"
	"google.golang.org/api/container/v1"
	"google.golang.org/api/dns/v1"
	"google.golang.org/api/storage/v1"
)

// Config is the configuration structure used to instantiate the Google
// provider.
type Config struct {
	AccountFile string
	Project     string
	Region      string

	clientCompute *compute.Service
	clientContainer *container.Service
	clientDns *dns.Service
	clientStorage *storage.Service
}

func (c *Config) loadAndValidate() error {
	var account accountFile

	if c.AccountFile == "" {
		c.AccountFile = os.Getenv("GOOGLE_ACCOUNT_FILE")
	}
	if c.Project == "" {
		c.Project = os.Getenv("GOOGLE_PROJECT")
	}
	if c.Region == "" {
		c.Region = os.Getenv("GOOGLE_REGION")
	}

	var client *http.Client

	if c.AccountFile != "" {
		contents := c.AccountFile

		// Assume account_file is a JSON string
		if err := parseJSON(&account, contents); err != nil {
			// If account_file was not JSON, assume it is a file path instead
			if _, err := os.Stat(c.AccountFile); os.IsNotExist(err) {
				return fmt.Errorf(
					"account_file path does not exist: %s",
					c.AccountFile)
			}

			b, err := ioutil.ReadFile(c.AccountFile)
			if err != nil {
				return fmt.Errorf(
					"Error reading account_file from path '%s': %s",
					c.AccountFile,
					err)
			}

			contents = string(b)

			if err := parseJSON(&account, contents); err != nil {
				return fmt.Errorf(
					"Error parsing account file '%s': %s",
					contents,
					err)
			}
		}

		clientScopes := []string{
			"https://www.googleapis.com/auth/compute",
			"https://www.googleapis.com/auth/cloud-platform",
			"https://www.googleapis.com/auth/ndev.clouddns.readwrite",
			"https://www.googleapis.com/auth/devstorage.full_control",
		}

		// Get the token for use in our requests
		log.Printf("[INFO] Requesting Google token...")
		log.Printf("[INFO]   -- Email: %s", account.ClientEmail)
		log.Printf("[INFO]   -- Scopes: %s", clientScopes)
		log.Printf("[INFO]   -- Private Key Length: %d", len(account.PrivateKey))

		conf := jwt.Config{
			Email:      account.ClientEmail,
			PrivateKey: []byte(account.PrivateKey),
			Scopes:     clientScopes,
			TokenURL:   "https://accounts.google.com/o/oauth2/token",
		}

		// Initiate an http.Client. The following GET request will be
		// authorized and authenticated on the behalf of
		// your service account.
		client = conf.Client(oauth2.NoContext)

	} else {
		log.Printf("[INFO] Requesting Google token via GCE Service Role...")
		client = &http.Client{
			Transport: &oauth2.Transport{
				// Fetch from Google Compute Engine's metadata server to retrieve
				// an access token for the provided account.
				// If no account is specified, "default" is used.
				Source: google.ComputeTokenSource(""),
			},
		}

	}

	// Build UserAgent
	versionString := "0.0.0"
	// TODO(dcunnin): Use Terraform's version code from version.go
	// versionString := main.Version
	// if main.VersionPrerelease != "" {
	// 	versionString = fmt.Sprintf("%s-%s", versionString, main.VersionPrerelease)
	// }
	userAgent := fmt.Sprintf(
		"(%s %s) Terraform/%s", runtime.GOOS, runtime.GOARCH, versionString)

	var err error

	log.Printf("[INFO] Instantiating GCE client...")
	c.clientCompute, err = compute.New(client)
	if err != nil {
		return err
	}
	c.clientCompute.UserAgent = userAgent

	log.Printf("[INFO] Instantiating GKE client...")
	c.clientContainer, err = container.New(client)
	if err != nil {
		return err
	}
	c.clientContainer.UserAgent = userAgent

	log.Printf("[INFO] Instantiating Google Cloud DNS client...")
	c.clientDns, err = dns.New(client)
	if err != nil {
		return err
	}
	c.clientDns.UserAgent = userAgent

	log.Printf("[INFO] Instantiating Google Storage Client...")
	c.clientStorage, err = storage.New(client)
	if err != nil {
		return err
	}
	c.clientStorage.UserAgent = userAgent

	return nil
}

// accountFile represents the structure of the account file JSON file.
type accountFile struct {
	PrivateKeyId string `json:"private_key_id"`
	PrivateKey   string `json:"private_key"`
	ClientEmail  string `json:"client_email"`
	ClientId     string `json:"client_id"`
}

func parseJSON(result interface{}, contents string) error {
	r := strings.NewReader(contents)
	dec := json.NewDecoder(r)

	return dec.Decode(result)
}
<|endoftext|>


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

// Package input contains the code processes new room events
package input

import (
	"context"
	"encoding/json"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/matrix-org/dendrite/internal/hooks"
	"github.com/matrix-org/dendrite/roomserver/acls"
	"github.com/matrix-org/dendrite/roomserver/api"
	"github.com/matrix-org/dendrite/roomserver/storage"
	"github.com/matrix-org/gomatrixserverlib"
	log "github.com/sirupsen/logrus"
	"go.uber.org/atomic"
)

type Inputer struct {
	DB                   storage.Database
	Producer             sarama.SyncProducer
	ServerName           gomatrixserverlib.ServerName
	ACLs                 *acls.ServerACLs
	OutputRoomEventTopic string

	workers sync.Map // room ID -> *inputWorker
}

type inputTask struct {
	ctx   context.Context
	event *api.InputRoomEvent
	wg    *sync.WaitGroup
	err   error // written back by worker, only safe to read when all tasks are done
}

type inputWorker struct {
	r       *Inputer
	running atomic.Bool
	input   chan *inputTask
}

func (w *inputWorker) start() {
	if !w.running.CAS(false, true) {
		return
	}
	defer w.running.Store(false)
	for {
		select {
		case task := <-w.input:
			hooks.Run(hooks.KindNewEventReceived, task.event.Event)
			_, task.err = w.r.processRoomEvent(task.ctx, task.event)
			if task.err == nil {
				hooks.Run(hooks.KindNewEventPersisted, task.event.Event)
			}
			task.wg.Done()
		case <-time.After(time.Second * 5):
			return
		}
	}
}

// WriteOutputEvents implements OutputRoomEventWriter
func (r *Inputer) WriteOutputEvents(roomID string, updates []api.OutputEvent) error {
	messages := make([]*sarama.ProducerMessage, len(updates))
	for i := range updates {
		value, err := json.Marshal(updates[i])
		if err != nil {
			return err
		}
		logger := log.WithFields(log.Fields{
			"room_id": roomID,
			"type":    updates[i].Type,
		})
		if updates[i].NewRoomEvent != nil {
			logger = logger.WithFields(log.Fields{
				"event_type":     updates[i].NewRoomEvent.Event.Type(),
				"event_id":       updates[i].NewRoomEvent.Event.EventID(),
				"adds_state":     len(updates[i].NewRoomEvent.AddsStateEventIDs),
				"removes_state":  len(updates[i].NewRoomEvent.RemovesStateEventIDs),
				"send_as_server": updates[i].NewRoomEvent.SendAsServer,
				"sender":         updates[i].NewRoomEvent.Event.Sender(),
			})
			if updates[i].NewRoomEvent.Event.Type() == "m.room.server_acl" && updates[i].NewRoomEvent.Event.StateKeyEquals("") {
				ev := updates[i].NewRoomEvent.Event.Unwrap()
				defer r.ACLs.OnServerACLUpdate(ev)
			}
		}
		logger.Infof("Producing to topic '%s'", r.OutputRoomEventTopic)
		messages[i] = &sarama.ProducerMessage{
			Topic: r.OutputRoomEventTopic,
			Key:   sarama.StringEncoder(roomID),
			Value: sarama.ByteEncoder(value),
		}
	}
	errs := r.Producer.SendMessages(messages)
	if errs != nil {
		for _, err := range errs.(sarama.ProducerErrors) {
			log.WithError(err).WithField("message_bytes", err.Msg.Value.Length()).Error("Write to kafka failed")
		}
	}
	return errs
}

// InputRoomEvents implements api.RoomserverInternalAPI
func (r *Inputer) InputRoomEvents(
	ctx context.Context,
	request *api.InputRoomEventsRequest,
	response *api.InputRoomEventsResponse,
) {
	// Create a wait group. Each task that we dispatch will call Done on
	// this wait group so that we know when all of our events have been
	// processed.
	wg := &sync.WaitGroup{}
	wg.Add(len(request.InputRoomEvents))
	tasks := make([]*inputTask, len(request.InputRoomEvents))

	for i, e := range request.InputRoomEvents {
		// Work out if we are running per-room workers or if we're just doing
		// it on a global basis (e.g. SQLite).
		roomID := "global"
		if r.DB.SupportsConcurrentRoomInputs() {
			roomID = e.Event.RoomID()
		}

		// Look up the worker, or create it if it doesn't exist. This channel
		// is buffered to reduce the chance that we'll be blocked by another
		// room - the channel will be quite small as it's just pointer types.
		w, _ := r.workers.LoadOrStore(roomID, &inputWorker{
			r:     r,
			input: make(chan *inputTask, 10),
		})
		worker := w.(*inputWorker)

		// Create a task. This contains the input event and a reference to
		// the wait group, so that the worker can notify us when this specific
		// task has been finished.
		tasks[i] = &inputTask{
			ctx:   ctx,
			event: &request.InputRoomEvents[i],
			wg:    wg,
		}

		// Send the task to the worker.
		go worker.start()
		worker.input <- tasks[i]
	}

	// Wait for all of the workers to return results about our tasks.
	wg.Wait()

	// If any of the tasks returned an error, we should probably report
	// that back to the caller.
	for _, task := range tasks {
		if task.err != nil {
			response.ErrMsg = task.err.Error()
			_, rejected := task.err.(*gomatrixserverlib.NotAllowed)
			response.NotAllowed = rejected
			return
		}
	}
}
<commit_msg>Don't start more goroutines than needed on RS input, increase input worker buffer size (#1638)<commit_after>// Copyright 2017 Vector Creations Ltd
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package input contains the code processes new room events
package input

import (
	"context"
	"encoding/json"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/matrix-org/dendrite/internal/hooks"
	"github.com/matrix-org/dendrite/roomserver/acls"
	"github.com/matrix-org/dendrite/roomserver/api"
	"github.com/matrix-org/dendrite/roomserver/storage"
	"github.com/matrix-org/gomatrixserverlib"
	log "github.com/sirupsen/logrus"
	"go.uber.org/atomic"
)

type Inputer struct {
	DB                   storage.Database
	Producer             sarama.SyncProducer
	ServerName           gomatrixserverlib.ServerName
	ACLs                 *acls.ServerACLs
	OutputRoomEventTopic string

	workers sync.Map // room ID -> *inputWorker
}

type inputTask struct {
	ctx   context.Context
	event *api.InputRoomEvent
	wg    *sync.WaitGroup
	err   error // written back by worker, only safe to read when all tasks are done
}

type inputWorker struct {
	r       *Inputer
	running atomic.Bool
	input   chan *inputTask
}

// Guarded by a CAS on w.running
func (w *inputWorker) start() {
	defer w.running.Store(false)
	for {
		select {
		case task := <-w.input:
			hooks.Run(hooks.KindNewEventReceived, task.event.Event)
			_, task.err = w.r.processRoomEvent(task.ctx, task.event)
			if task.err == nil {
				hooks.Run(hooks.KindNewEventPersisted, task.event.Event)
			}
			task.wg.Done()
		case <-time.After(time.Second * 5):
			return
		}
	}
}

// WriteOutputEvents implements OutputRoomEventWriter
func (r *Inputer) WriteOutputEvents(roomID string, updates []api.OutputEvent) error {
	messages := make([]*sarama.ProducerMessage, len(updates))
	for i := range updates {
		value, err := json.Marshal(updates[i])
		if err != nil {
			return err
		}
		logger := log.WithFields(log.Fields{
			"room_id": roomID,
			"type":    updates[i].Type,
		})
		if updates[i].NewRoomEvent != nil {
			logger = logger.WithFields(log.Fields{
				"event_type":     updates[i].NewRoomEvent.Event.Type(),
				"event_id":       updates[i].NewRoomEvent.Event.EventID(),
				"adds_state":     len(updates[i].NewRoomEvent.AddsStateEventIDs),
				"removes_state":  len(updates[i].NewRoomEvent.RemovesStateEventIDs),
				"send_as_server": updates[i].NewRoomEvent.SendAsServer,
				"sender":         updates[i].NewRoomEvent.Event.Sender(),
			})
			if updates[i].NewRoomEvent.Event.Type() == "m.room.server_acl" && updates[i].NewRoomEvent.Event.StateKeyEquals("") {
				ev := updates[i].NewRoomEvent.Event.Unwrap()
				defer r.ACLs.OnServerACLUpdate(ev)
			}
		}
		logger.Infof("Producing to topic '%s'", r.OutputRoomEventTopic)
		messages[i] = &sarama.ProducerMessage{
			Topic: r.OutputRoomEventTopic,
			Key:   sarama.StringEncoder(roomID),
			Value: sarama.ByteEncoder(value),
		}
	}
	errs := r.Producer.SendMessages(messages)
	if errs != nil {
		for _, err := range errs.(sarama.ProducerErrors) {
			log.WithError(err).WithField("message_bytes", err.Msg.Value.Length()).Error("Write to kafka failed")
		}
	}
	return errs
}

// InputRoomEvents implements api.RoomserverInternalAPI
func (r *Inputer) InputRoomEvents(
	ctx context.Context,
	request *api.InputRoomEventsRequest,
	response *api.InputRoomEventsResponse,
) {
	// Create a wait group. Each task that we dispatch will call Done on
	// this wait group so that we know when all of our events have been
	// processed.
	wg := &sync.WaitGroup{}
	wg.Add(len(request.InputRoomEvents))
	tasks := make([]*inputTask, len(request.InputRoomEvents))

	for i, e := range request.InputRoomEvents {
		// Work out if we are running per-room workers or if we're just doing
		// it on a global basis (e.g. SQLite).
		roomID := "global"
		if r.DB.SupportsConcurrentRoomInputs() {
			roomID = e.Event.RoomID()
		}

		// Look up the worker, or create it if it doesn't exist. This channel
		// is buffered to reduce the chance that we'll be blocked by another
		// room - the channel will be quite small as it's just pointer types.
		w, _ := r.workers.LoadOrStore(roomID, &inputWorker{
			r:     r,
			input: make(chan *inputTask, 32),
		})
		worker := w.(*inputWorker)

		// Create a task. This contains the input event and a reference to
		// the wait group, so that the worker can notify us when this specific
		// task has been finished.
		tasks[i] = &inputTask{
			ctx:   ctx,
			event: &request.InputRoomEvents[i],
			wg:    wg,
		}

		// Send the task to the worker.
		if worker.running.CAS(false, true) {
			go worker.start()
		}
		worker.input <- tasks[i]
	}

	// Wait for all of the workers to return results about our tasks.
	wg.Wait()

	// If any of the tasks returned an error, we should probably report
	// that back to the caller.
	for _, task := range tasks {
		if task.err != nil {
			response.ErrMsg = task.err.Error()
			_, rejected := task.err.(*gomatrixserverlib.NotAllowed)
			response.NotAllowed = rejected
			return
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build go1.9

package ipv6

import (
	"net"
	"runtime"
	"syscall"

	"golang.org/x/net/internal/socket"
)

// BUG(mikio): On Windows, the ReadBatch and WriteBatch methods of
// PacketConn are not implemented.

// A Message represents an IO message.
//
//	type Message struct {
//		Buffers [][]byte
//		OOB     []byte
//		Addr    net.Addr
//		N       int
//		NN      int
//		Flags   int
//	}
//
// The Buffers fields represents a list of contiguous buffers, which
// can be used for vectored IO, for example, putting a header and a
// payload in each slice.
// When writing, the Buffers field must contain at least one byte to
// write.
// When reading, the Buffers field will always contain a byte to read.
//
// The OOB field contains protocol-specific control or miscellaneous
// ancillary data known as out-of-band data.
// It can be nil when not required.
//
// The Addr field specifies a destination address when writing.
// It can be nil when the underlying protocol of the endpoint uses
// connection-oriented communication.
// After a successful read, it may contain the source address on the
// received packet.
//
// The N field indicates the number of bytes read or written from/to
// Buffers.
//
// The NN field indicates the number of bytes read or written from/to
// OOB.
//
// The Flags field contains protocol-specific information on the
// received message.
type Message = socket.Message

// ReadBatch reads a batch of messages.
//
// The provided flags is a set of platform-dependent flags, such as
// syscall.MSG_PEEK.
//
// On a successful read it returns the number of messages received, up
// to len(ms).
//
// On Linux, a batch read will be optimized.
// On other platforms, this method will read only a single message.
func (c *payloadHandler) ReadBatch(ms []Message, flags int) (int, error) {
	if !c.ok() {
		return 0, syscall.EINVAL
	}
	switch runtime.GOOS {
	case "linux":
		n, err := c.RecvMsgs([]socket.Message(ms), flags)
		if err != nil {
			err = &net.OpError{Op: "read", Net: c.PacketConn.LocalAddr().Network(), Source: c.PacketConn.LocalAddr(), Err: err}
		}
		return n, err
	default:
		n := 1
		err := c.RecvMsg(&ms[0], flags)
		if err != nil {
			n = 0
			err = &net.OpError{Op: "read", Net: c.PacketConn.LocalAddr().Network(), Source: c.PacketConn.LocalAddr(), Err: err}
		}
		return n, err
	}
}

// WriteBatch writes a batch of messages.
//
// The provided flags is a set of platform-dependent flags, such as
// syscall.MSG_DONTROUTE.
//
// It returns the number of messages written on a successful write.
//
// On Linux, a batch write will be optimized.
// On other platforms, this method will write only a single message.
func (c *payloadHandler) WriteBatch(ms []Message, flags int) (int, error) {
	if !c.ok() {
		return 0, syscall.EINVAL
	}
	switch runtime.GOOS {
	case "linux":
		n, err := c.SendMsgs([]socket.Message(ms), flags)
		if err != nil {
			err = &net.OpError{Op: "write", Net: c.PacketConn.LocalAddr().Network(), Source: c.PacketConn.LocalAddr(), Err: err}
		}
		return n, err
	default:
		n := 1
		err := c.SendMsg(&ms[0], flags)
		if err != nil {
			n = 0
			err = &net.OpError{Op: "write", Net: c.PacketConn.LocalAddr().Network(), Source: c.PacketConn.LocalAddr(), Err: err}
		}
		return n, err
	}
}
<commit_msg>ipv6: enable {Read,Write}Batch methods of PacketConn on NetBSD<commit_after>// Copyright 2017 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build go1.9

package ipv6

import (
	"net"
	"runtime"
	"syscall"

	"golang.org/x/net/internal/socket"
)

// BUG(mikio): On Windows, the ReadBatch and WriteBatch methods of
// PacketConn are not implemented.

// A Message represents an IO message.
//
//	type Message struct {
//		Buffers [][]byte
//		OOB     []byte
//		Addr    net.Addr
//		N       int
//		NN      int
//		Flags   int
//	}
//
// The Buffers fields represents a list of contiguous buffers, which
// can be used for vectored IO, for example, putting a header and a
// payload in each slice.
// When writing, the Buffers field must contain at least one byte to
// write.
// When reading, the Buffers field will always contain a byte to read.
//
// The OOB field contains protocol-specific control or miscellaneous
// ancillary data known as out-of-band data.
// It can be nil when not required.
//
// The Addr field specifies a destination address when writing.
// It can be nil when the underlying protocol of the endpoint uses
// connection-oriented communication.
// After a successful read, it may contain the source address on the
// received packet.
//
// The N field indicates the number of bytes read or written from/to
// Buffers.
//
// The NN field indicates the number of bytes read or written from/to
// OOB.
//
// The Flags field contains protocol-specific information on the
// received message.
type Message = socket.Message

// ReadBatch reads a batch of messages.
//
// The provided flags is a set of platform-dependent flags, such as
// syscall.MSG_PEEK.
//
// On a successful read it returns the number of messages received, up
// to len(ms).
//
// On Linux and NetBSD, a batch read will be optimized.
// On other platforms, this method will read only a single message.
func (c *payloadHandler) ReadBatch(ms []Message, flags int) (int, error) {
	if !c.ok() {
		return 0, syscall.EINVAL
	}
	switch runtime.GOOS {
	case "linux", "netbsd":
		n, err := c.RecvMsgs([]socket.Message(ms), flags)
		if err != nil {
			err = &net.OpError{Op: "read", Net: c.PacketConn.LocalAddr().Network(), Source: c.PacketConn.LocalAddr(), Err: err}
		}
		return n, err
	default:
		n := 1
		err := c.RecvMsg(&ms[0], flags)
		if err != nil {
			n = 0
			err = &net.OpError{Op: "read", Net: c.PacketConn.LocalAddr().Network(), Source: c.PacketConn.LocalAddr(), Err: err}
		}
		return n, err
	}
}

// WriteBatch writes a batch of messages.
//
// The provided flags is a set of platform-dependent flags, such as
// syscall.MSG_DONTROUTE.
//
// It returns the number of messages written on a successful write.
//
// On Linux and NetBSD, a batch write will be optimized.
// On other platforms, this method will write only a single message.
func (c *payloadHandler) WriteBatch(ms []Message, flags int) (int, error) {
	if !c.ok() {
		return 0, syscall.EINVAL
	}
	switch runtime.GOOS {
	case "linux", "netbsd":
		n, err := c.SendMsgs([]socket.Message(ms), flags)
		if err != nil {
			err = &net.OpError{Op: "write", Net: c.PacketConn.LocalAddr().Network(), Source: c.PacketConn.LocalAddr(), Err: err}
		}
		return n, err
	default:
		n := 1
		err := c.SendMsg(&ms[0], flags)
		if err != nil {
			n = 0
			err = &net.OpError{Op: "write", Net: c.PacketConn.LocalAddr().Network(), Source: c.PacketConn.LocalAddr(), Err: err}
		}
		return n, err
	}
}
<|endoftext|>


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

package perform

import (
	"context"
	"fmt"
	"strings"

	fsAPI "github.com/matrix-org/dendrite/federationsender/api"
	"github.com/matrix-org/dendrite/roomserver/api"
	"github.com/matrix-org/dendrite/roomserver/internal/helpers"
	"github.com/matrix-org/dendrite/roomserver/internal/input"
	"github.com/matrix-org/dendrite/roomserver/storage"
	"github.com/matrix-org/dendrite/setup/config"
	"github.com/matrix-org/gomatrixserverlib"
)

type Leaver struct {
	Cfg   *config.RoomServer
	DB    storage.Database
	FSAPI fsAPI.FederationSenderInternalAPI

	Inputer *input.Inputer
}

// WriteOutputEvents implements OutputRoomEventWriter
func (r *Leaver) PerformLeave(
	ctx context.Context,
	req *api.PerformLeaveRequest,
	res *api.PerformLeaveResponse,
) ([]api.OutputEvent, error) {
	_, domain, err := gomatrixserverlib.SplitID('@', req.UserID)
	if err != nil {
		return nil, fmt.Errorf("Supplied user ID %q in incorrect format", req.UserID)
	}
	if domain != r.Cfg.Matrix.ServerName {
		return nil, fmt.Errorf("User %q does not belong to this homeserver", req.UserID)
	}
	if strings.HasPrefix(req.RoomID, "!") {
		return r.performLeaveRoomByID(ctx, req, res)
	}
	return nil, fmt.Errorf("Room ID %q is invalid", req.RoomID)
}

func (r *Leaver) performLeaveRoomByID(
	ctx context.Context,
	req *api.PerformLeaveRequest,
	res *api.PerformLeaveResponse, // nolint:unparam
) ([]api.OutputEvent, error) {
	// If there's an invite outstanding for the room then respond to
	// that.
	isInvitePending, senderUser, eventID, err := helpers.IsInvitePending(ctx, r.DB, req.RoomID, req.UserID)
	if err == nil && isInvitePending {
		return r.performRejectInvite(ctx, req, res, senderUser, eventID)
	}

	// There's no invite pending, so first of all we want to find out
	// if the room exists and if the user is actually in it.
	latestReq := api.QueryLatestEventsAndStateRequest{
		RoomID: req.RoomID,
		StateToFetch: []gomatrixserverlib.StateKeyTuple{
			{
				EventType: gomatrixserverlib.MRoomMember,
				StateKey:  req.UserID,
			},
		},
	}
	latestRes := api.QueryLatestEventsAndStateResponse{}
	if err = helpers.QueryLatestEventsAndState(ctx, r.DB, &latestReq, &latestRes); err != nil {
		return nil, err
	}
	if !latestRes.RoomExists {
		return nil, fmt.Errorf("Room %q does not exist", req.RoomID)
	}

	// Now let's see if the user is in the room.
	if len(latestRes.StateEvents) == 0 {
		return nil, fmt.Errorf("User %q is not a member of room %q", req.UserID, req.RoomID)
	}
	membership, err := latestRes.StateEvents[0].Membership()
	if err != nil {
		return nil, fmt.Errorf("Error getting membership: %w", err)
	}
	if membership != gomatrixserverlib.Join {
		// TODO: should be able to handle "invite" in this case too, if
		// it's a case of kicking or banning or such
		return nil, fmt.Errorf("User %q is not joined to the room (membership is %q)", req.UserID, membership)
	}

	// Prepare the template for the leave event.
	userID := req.UserID
	eb := gomatrixserverlib.EventBuilder{
		Type:     gomatrixserverlib.MRoomMember,
		Sender:   userID,
		StateKey: &userID,
		RoomID:   req.RoomID,
		Redacts:  "",
	}
	if err = eb.SetContent(map[string]interface{}{"membership": "leave"}); err != nil {
		return nil, fmt.Errorf("eb.SetContent: %w", err)
	}
	if err = eb.SetUnsigned(struct{}{}); err != nil {
		return nil, fmt.Errorf("eb.SetUnsigned: %w", err)
	}

	// We know that the user is in the room at this point so let's build
	// a leave event.
	// TODO: Check what happens if the room exists on the server
	// but everyone has since left. I suspect it does the wrong thing.
	event, buildRes, err := buildEvent(ctx, r.DB, r.Cfg.Matrix, &eb)
	if err != nil {
		return nil, fmt.Errorf("eventutil.BuildEvent: %w", err)
	}

	// Give our leave event to the roomserver input stream. The
	// roomserver will process the membership change and notify
	// downstream automatically.
	inputReq := api.InputRoomEventsRequest{
		InputRoomEvents: []api.InputRoomEvent{
			{
				Kind:         api.KindNew,
				Event:        event.Headered(buildRes.RoomVersion),
				AuthEventIDs: event.AuthEventIDs(),
				SendAsServer: string(r.Cfg.Matrix.ServerName),
			},
		},
	}
	inputRes := api.InputRoomEventsResponse{}
	r.Inputer.InputRoomEvents(ctx, &inputReq, &inputRes)
	if err = inputRes.Err(); err != nil {
		return nil, fmt.Errorf("r.InputRoomEvents: %w", err)
	}

	return nil, nil
}

func (r *Leaver) performRejectInvite(
	ctx context.Context,
	req *api.PerformLeaveRequest,
	res *api.PerformLeaveResponse, // nolint:unparam
	senderUser, eventID string,
) ([]api.OutputEvent, error) {
	_, domain, err := gomatrixserverlib.SplitID('@', senderUser)
	if err != nil {
		return nil, fmt.Errorf("User ID %q invalid: %w", senderUser, err)
	}

	// Ask the federation sender to perform a federated leave for us.
	leaveReq := fsAPI.PerformLeaveRequest{
		RoomID:      req.RoomID,
		UserID:      req.UserID,
		ServerNames: []gomatrixserverlib.ServerName{domain},
	}
	leaveRes := fsAPI.PerformLeaveResponse{}
	if err := r.FSAPI.PerformLeave(ctx, &leaveReq, &leaveRes); err != nil {
		return nil, err
	}

	// Withdraw the invite, so that the sync API etc are
	// notified that we rejected it.
	return []api.OutputEvent{
		{
			Type: api.OutputTypeRetireInviteEvent,
			RetireInviteEvent: &api.OutputRetireInviteEvent{
				EventID:      eventID,
				Membership:   "leave",
				TargetUserID: req.UserID,
			},
		},
	}, nil
}
<commit_msg>Fix bug when rejecting invites (#1907)<commit_after>// Copyright 2020 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package perform

import (
	"context"
	"fmt"
	"strings"

	fsAPI "github.com/matrix-org/dendrite/federationsender/api"
	"github.com/matrix-org/dendrite/roomserver/api"
	"github.com/matrix-org/dendrite/roomserver/internal/helpers"
	"github.com/matrix-org/dendrite/roomserver/internal/input"
	"github.com/matrix-org/dendrite/roomserver/storage"
	"github.com/matrix-org/dendrite/setup/config"
	"github.com/matrix-org/gomatrixserverlib"
)

type Leaver struct {
	Cfg   *config.RoomServer
	DB    storage.Database
	FSAPI fsAPI.FederationSenderInternalAPI

	Inputer *input.Inputer
}

// WriteOutputEvents implements OutputRoomEventWriter
func (r *Leaver) PerformLeave(
	ctx context.Context,
	req *api.PerformLeaveRequest,
	res *api.PerformLeaveResponse,
) ([]api.OutputEvent, error) {
	_, domain, err := gomatrixserverlib.SplitID('@', req.UserID)
	if err != nil {
		return nil, fmt.Errorf("Supplied user ID %q in incorrect format", req.UserID)
	}
	if domain != r.Cfg.Matrix.ServerName {
		return nil, fmt.Errorf("User %q does not belong to this homeserver", req.UserID)
	}
	if strings.HasPrefix(req.RoomID, "!") {
		return r.performLeaveRoomByID(ctx, req, res)
	}
	return nil, fmt.Errorf("Room ID %q is invalid", req.RoomID)
}

func (r *Leaver) performLeaveRoomByID(
	ctx context.Context,
	req *api.PerformLeaveRequest,
	res *api.PerformLeaveResponse, // nolint:unparam
) ([]api.OutputEvent, error) {
	// If there's an invite outstanding for the room then respond to
	// that.
	isInvitePending, senderUser, eventID, err := helpers.IsInvitePending(ctx, r.DB, req.RoomID, req.UserID)
	if err == nil && isInvitePending {
		var host gomatrixserverlib.ServerName
		_, host, err = gomatrixserverlib.SplitID('@', senderUser)
		if err != nil {
			return nil, fmt.Errorf("Sender %q is invalid", senderUser)
		}
		if host != r.Cfg.Matrix.ServerName {
			return r.performFederatedRejectInvite(ctx, req, res, senderUser, eventID)
		}
	}

	// There's no invite pending, so first of all we want to find out
	// if the room exists and if the user is actually in it.
	latestReq := api.QueryLatestEventsAndStateRequest{
		RoomID: req.RoomID,
		StateToFetch: []gomatrixserverlib.StateKeyTuple{
			{
				EventType: gomatrixserverlib.MRoomMember,
				StateKey:  req.UserID,
			},
		},
	}
	latestRes := api.QueryLatestEventsAndStateResponse{}
	if err = helpers.QueryLatestEventsAndState(ctx, r.DB, &latestReq, &latestRes); err != nil {
		return nil, err
	}
	if !latestRes.RoomExists {
		return nil, fmt.Errorf("Room %q does not exist", req.RoomID)
	}

	// Now let's see if the user is in the room.
	if len(latestRes.StateEvents) == 0 {
		return nil, fmt.Errorf("User %q is not a member of room %q", req.UserID, req.RoomID)
	}
	membership, err := latestRes.StateEvents[0].Membership()
	if err != nil {
		return nil, fmt.Errorf("Error getting membership: %w", err)
	}
	if membership != gomatrixserverlib.Join && membership != gomatrixserverlib.Invite {
		return nil, fmt.Errorf("User %q is not joined to the room (membership is %q)", req.UserID, membership)
	}

	// Prepare the template for the leave event.
	userID := req.UserID
	eb := gomatrixserverlib.EventBuilder{
		Type:     gomatrixserverlib.MRoomMember,
		Sender:   userID,
		StateKey: &userID,
		RoomID:   req.RoomID,
		Redacts:  "",
	}
	if err = eb.SetContent(map[string]interface{}{"membership": "leave"}); err != nil {
		return nil, fmt.Errorf("eb.SetContent: %w", err)
	}
	if err = eb.SetUnsigned(struct{}{}); err != nil {
		return nil, fmt.Errorf("eb.SetUnsigned: %w", err)
	}

	// We know that the user is in the room at this point so let's build
	// a leave event.
	// TODO: Check what happens if the room exists on the server
	// but everyone has since left. I suspect it does the wrong thing.
	event, buildRes, err := buildEvent(ctx, r.DB, r.Cfg.Matrix, &eb)
	if err != nil {
		return nil, fmt.Errorf("eventutil.BuildEvent: %w", err)
	}

	// Give our leave event to the roomserver input stream. The
	// roomserver will process the membership change and notify
	// downstream automatically.
	inputReq := api.InputRoomEventsRequest{
		InputRoomEvents: []api.InputRoomEvent{
			{
				Kind:         api.KindNew,
				Event:        event.Headered(buildRes.RoomVersion),
				AuthEventIDs: event.AuthEventIDs(),
				SendAsServer: string(r.Cfg.Matrix.ServerName),
			},
		},
	}
	inputRes := api.InputRoomEventsResponse{}
	r.Inputer.InputRoomEvents(ctx, &inputReq, &inputRes)
	if err = inputRes.Err(); err != nil {
		return nil, fmt.Errorf("r.InputRoomEvents: %w", err)
	}

	return nil, nil
}

func (r *Leaver) performFederatedRejectInvite(
	ctx context.Context,
	req *api.PerformLeaveRequest,
	res *api.PerformLeaveResponse, // nolint:unparam
	senderUser, eventID string,
) ([]api.OutputEvent, error) {
	_, domain, err := gomatrixserverlib.SplitID('@', senderUser)
	if err != nil {
		return nil, fmt.Errorf("User ID %q invalid: %w", senderUser, err)
	}

	// Ask the federation sender to perform a federated leave for us.
	leaveReq := fsAPI.PerformLeaveRequest{
		RoomID:      req.RoomID,
		UserID:      req.UserID,
		ServerNames: []gomatrixserverlib.ServerName{domain},
	}
	leaveRes := fsAPI.PerformLeaveResponse{}
	if err := r.FSAPI.PerformLeave(ctx, &leaveReq, &leaveRes); err != nil {
		return nil, err
	}

	// Withdraw the invite, so that the sync API etc are
	// notified that we rejected it.
	return []api.OutputEvent{
		{
			Type: api.OutputTypeRetireInviteEvent,
			RetireInviteEvent: &api.OutputRetireInviteEvent{
				EventID:      eventID,
				Membership:   "leave",
				TargetUserID: req.UserID,
			},
		},
	}, nil
}
<|endoftext|>


<commit_before>// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see <http://www.gnu.org/licenses/>.

package irc

import (
	"fmt"
	"net"
	"strings"
)

type Hook func(*Client, Message) error

type Client struct {
	conn     net.Conn
	hooks    map[string][]Hook
	Nickname string
	Realname string
	Channels []string
}

func NewClient(conn net.Conn) *Client {
	c := &Client{
		conn:  conn,
		hooks: make(map[string][]Hook),
	}

	c.CmdHook("join", c.joinCmd)
	c.CmdHook("part", c.partCmd)
	c.CmdHook("kick", c.kickCmd)

	c.CmdHook("ping", c.pingCmd)
	return c
}

func (c *Client) Setup(nick, name, host string) {
	c.Nickname = nick
	c.Realname = name

	c.Write("USER %s %s * :%s", c.Nickname, host, c.Realname)
	c.Write("NICK %s", c.Nickname)
}

func (c *Client) Write(format string, argv ...interface{}) error {
	_, err := fmt.Fprintf(c.conn, "%s\r\n", fmt.Sprintf(format, argv...))
	if err != nil {
		return err
	}

	return nil
}

func (c *Client) Handle(data string, ch chan error) {
	msg := parseMessage(data)
	hooks, ok := c.hooks[msg.Command]
	if ok {
		for _, hook := range hooks {
			go func(h Hook) {
				if err := h(c, msg); err != nil {
					ch <- err
				}
			}(hook)
		}
	}
}

func (c *Client) CmdHook(cmd string, hook Hook) {
	c.hooks[cmd] = append(c.hooks[cmd], hook)
}

func (c *Client) Connected(channel string) bool {
	for _, c := range c.Channels {
		if c == channel {
			return true
		}
	}

	return false
}

func (c *Client) remove(channel string) {
	var newChans []string
	for _, c := range c.Channels {
		if c != channel {
			newChans = append(newChans, c)
		}
	}

	c.Channels = newChans
}

func (c *Client) joinCmd(client *Client, msg Message) error {
	if msg.Sender.Name == c.Nickname {
		c.Channels = append(c.Channels, msg.Data)
	}

	return nil
}

func (c *Client) partCmd(client *Client, msg Message) error {
	if msg.Sender.Name == c.Nickname {
		c.remove(msg.Data)
	}

	return nil
}

func (c *Client) kickCmd(client *Client, msg Message) error {
	if msg.Data == c.Nickname {
		c.remove(strings.Fields(msg.Receiver)[0])
	}

	return nil
}

func (c *Client) pingCmd(client *Client, msg Message) error {
	return c.Write("PONG %s", msg.Data)
}
<commit_msg>irc: minimal hook cleanup<commit_after>// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see <http://www.gnu.org/licenses/>.

package irc

import (
	"fmt"
	"net"
	"strings"
)

type Hook func(*Client, Message) error

type Client struct {
	conn     net.Conn
	hooks    map[string][]Hook
	Nickname string
	Realname string
	Channels []string
}

func NewClient(conn net.Conn) *Client {
	c := &Client{
		conn:  conn,
		hooks: make(map[string][]Hook),
	}

	c.CmdHook("join", joinCmd)
	c.CmdHook("part", partCmd)
	c.CmdHook("kick", kickCmd)

	c.CmdHook("ping", pingCmd)
	return c
}

func (c *Client) Setup(nick, name, host string) {
	c.Nickname = nick
	c.Realname = name

	c.Write("USER %s %s * :%s", c.Nickname, host, c.Realname)
	c.Write("NICK %s", c.Nickname)
}

func (c *Client) Connected(channel string) bool {
	for _, c := range c.Channels {
		if c == channel {
			return true
		}
	}

	return false
}

func (c *Client) Write(format string, argv ...interface{}) error {
	_, err := fmt.Fprintf(c.conn, "%s\r\n", fmt.Sprintf(format, argv...))
	if err != nil {
		return err
	}

	return nil
}

func (c *Client) Handle(data string, ch chan error) {
	msg := parseMessage(data)
	hooks, ok := c.hooks[msg.Command]
	if ok {
		for _, hook := range hooks {
			go func(h Hook) {
				if err := h(c, msg); err != nil {
					ch <- err
				}
			}(hook)
		}
	}
}

func (c *Client) CmdHook(cmd string, hook Hook) {
	c.hooks[cmd] = append(c.hooks[cmd], hook)
}

func joinCmd(client *Client, msg Message) error {
	if msg.Sender.Name == client.Nickname {
		client.Channels = append(client.Channels, msg.Data)
	}

	return nil
}

func partCmd(client *Client, msg Message) error {
	if msg.Sender.Name == client.Nickname {
		client.Channels = remove(msg.Data, client.Channels)
	}

	return nil
}

func kickCmd(client *Client, msg Message) error {
	if msg.Data == client.Nickname {
		target := strings.Fields(msg.Receiver)[0]
		client.Channels = remove(target, client.Channels)
	}

	return nil
}

func pingCmd(client *Client, msg Message) error {
	return client.Write("PONG %s", msg.Data)
}

// remove deletes a given element from a given slice. A new slice
// which does not contain the given element is returned.
func remove(element string, slice []string) []string {
	var newSlice []string
	for _, e := range slice {
		if e != element {
			newSlice = append(newSlice, e)
		}
	}

	return newSlice
}
<|endoftext|>


<commit_before>package proxyproto

import (
	"bytes"
	"io"
	"net"
	"testing"
)

func TestPassthrough(t *testing.T) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	pl := &Listener{Listener: l}
	defer pl.Close()

	go func() {
		conn, err := net.Dial("tcp", pl.Addr().String())
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		defer conn.Close()

		conn.Write([]byte("ping"))
		recv := make([]byte, 4)
		_, err = conn.Read(recv)
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		if !bytes.Equal(recv, []byte("pong")) {
			t.Fatalf("bad: %v", recv)
		}
	}()

	conn, err := pl.Accept()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer conn.Close()

	recv := make([]byte, 4)
	_, err = conn.Read(recv)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if !bytes.Equal(recv, []byte("ping")) {
		t.Fatalf("bad: %v", recv)
	}

	if _, err := conn.Write([]byte("pong")); err != nil {
		t.Fatalf("err: %v", err)
	}
}

func TestParse_ipv4(t *testing.T) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	pl := &Listener{Listener: l}
	defer pl.Close()

	go func() {
		conn, err := net.Dial("tcp", pl.Addr().String())
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		defer conn.Close()

		// Write out the header!
		header := "PROXY TCP4 10.1.1.1 20.2.2.2 1000 2000\r\nfoo"
		conn.Write([]byte(header))

		conn.Write([]byte("ping"))
		recv := make([]byte, 4)
		_, err = conn.Read(recv)
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		if !bytes.Equal(recv, []byte("pong")) {
			t.Fatalf("bad: %v", recv)
		}
	}()

	conn, err := pl.Accept()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer conn.Close()

	recv := make([]byte, 7)
	_, err = conn.Read(recv)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if !bytes.Equal(recv, []byte("fooping")) {
		t.Fatalf("bad: %v", recv)
	}

	if _, err := conn.Write([]byte("pong")); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Check the remote addr
	addr := conn.RemoteAddr().(*net.TCPAddr)
	if addr.IP.String() != "10.1.1.1" {
		t.Fatalf("bad: %v", addr)
	}
	if addr.Port != 1000 {
		t.Fatalf("bad: %v", addr)
	}
}

func TestParse_ipv6(t *testing.T) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	pl := &Listener{Listener: l}
	defer pl.Close()

	go func() {
		conn, err := net.Dial("tcp", pl.Addr().String())
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		defer conn.Close()

		// Write out the header!
		header := "PROXY TCP6 ffff::ffff ffff::ffff 1000 2000\r\n"
		conn.Write([]byte(header))

		conn.Write([]byte("ping"))
		recv := make([]byte, 4)
		_, err = conn.Read(recv)
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		if !bytes.Equal(recv, []byte("pong")) {
			t.Fatalf("bad: %v", recv)
		}
	}()

	conn, err := pl.Accept()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer conn.Close()

	recv := make([]byte, 4)
	_, err = conn.Read(recv)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if !bytes.Equal(recv, []byte("ping")) {
		t.Fatalf("bad: %v", recv)
	}

	if _, err := conn.Write([]byte("pong")); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Check the remote addr
	addr := conn.RemoteAddr().(*net.TCPAddr)
	if addr.IP.String() != "ffff::ffff" {
		t.Fatalf("bad: %v", addr)
	}
	if addr.Port != 1000 {
		t.Fatalf("bad: %v", addr)
	}
}

func TestParse_BadHeader(t *testing.T) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	pl := &Listener{Listener: l}
	defer pl.Close()

	go func() {
		conn, err := net.Dial("tcp", pl.Addr().String())
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		defer conn.Close()

		// Write out the header!
		header := "PROXY TCP4 what 127.0.0.1 1000 2000\r\n"
		conn.Write([]byte(header))

		conn.Write([]byte("ping"))

		recv := make([]byte, 4)
		_, err = conn.Read(recv)
		if err != io.EOF {
			t.Fatalf("err: %v", err)
		}
	}()

	conn, err := pl.Accept()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer conn.Close()

	// Check the remote addr, should be the local addr
	addr := conn.RemoteAddr().(*net.TCPAddr)
	if addr.IP.String() != "127.0.0.1" {
		t.Fatalf("bad: %v", addr)
	}

	// Read should fail
	recv := make([]byte, 4)
	_, err = conn.Read(recv)
	if err == nil {
		t.Fatalf("err: %v", err)
	}
}
<commit_msg>router/proxyproto: Don't call t.Fatalf from goroutines<commit_after>package proxyproto

import (
	"bytes"
	"net"
	"testing"
)

func TestPassthrough(t *testing.T) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	pl := &Listener{Listener: l}
	defer pl.Close()

	go func() {
		conn, err := net.Dial("tcp", pl.Addr().String())
		if err != nil {
			t.Errorf("err: %v", err)
			return
		}
		defer conn.Close()

		conn.Write([]byte("ping"))
		recv := make([]byte, 4)
		_, err = conn.Read(recv)
		if err != nil {
			t.Errorf("err: %v", err)
			return
		}
		if !bytes.Equal(recv, []byte("pong")) {
			t.Errorf("bad: %v", recv)
			return
		}
	}()

	conn, err := pl.Accept()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer conn.Close()

	recv := make([]byte, 4)
	_, err = conn.Read(recv)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if !bytes.Equal(recv, []byte("ping")) {
		t.Fatalf("bad: %v", recv)
	}

	if _, err := conn.Write([]byte("pong")); err != nil {
		t.Fatalf("err: %v", err)
	}
}

func TestParse_ipv4(t *testing.T) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	pl := &Listener{Listener: l}
	defer pl.Close()

	go func() {
		conn, err := net.Dial("tcp", pl.Addr().String())
		if err != nil {
			t.Errorf("err: %v", err)
			return
		}
		defer conn.Close()

		// Write out the header!
		header := "PROXY TCP4 10.1.1.1 20.2.2.2 1000 2000\r\nfoo"
		conn.Write([]byte(header))

		conn.Write([]byte("ping"))
		recv := make([]byte, 4)
		_, err = conn.Read(recv)
		if err != nil {
			t.Errorf("err: %v", err)
			return
		}
		if !bytes.Equal(recv, []byte("pong")) {
			t.Errorf("bad: %v", recv)
			return
		}
	}()

	conn, err := pl.Accept()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer conn.Close()

	recv := make([]byte, 7)
	_, err = conn.Read(recv)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if !bytes.Equal(recv, []byte("fooping")) {
		t.Fatalf("bad: %v", recv)
	}

	if _, err := conn.Write([]byte("pong")); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Check the remote addr
	addr := conn.RemoteAddr().(*net.TCPAddr)
	if addr.IP.String() != "10.1.1.1" {
		t.Fatalf("bad: %v", addr)
	}
	if addr.Port != 1000 {
		t.Fatalf("bad: %v", addr)
	}
}

func TestParse_ipv6(t *testing.T) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	pl := &Listener{Listener: l}
	defer pl.Close()

	go func() {
		conn, err := net.Dial("tcp", pl.Addr().String())
		if err != nil {
			t.Errorf("err: %v", err)
			return
		}
		defer conn.Close()

		// Write out the header!
		header := "PROXY TCP6 ffff::ffff ffff::ffff 1000 2000\r\n"
		conn.Write([]byte(header))

		conn.Write([]byte("ping"))
		recv := make([]byte, 4)
		_, err = conn.Read(recv)
		if err != nil {
			t.Errorf("err: %v", err)
			return
		}
		if !bytes.Equal(recv, []byte("pong")) {
			t.Errorf("bad: %v", recv)
			return
		}
	}()

	conn, err := pl.Accept()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer conn.Close()

	recv := make([]byte, 4)
	_, err = conn.Read(recv)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if !bytes.Equal(recv, []byte("ping")) {
		t.Fatalf("bad: %v", recv)
	}

	if _, err := conn.Write([]byte("pong")); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Check the remote addr
	addr := conn.RemoteAddr().(*net.TCPAddr)
	if addr.IP.String() != "ffff::ffff" {
		t.Fatalf("bad: %v", addr)
	}
	if addr.Port != 1000 {
		t.Fatalf("bad: %v", addr)
	}
}

func TestParse_BadHeader(t *testing.T) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	pl := &Listener{Listener: l}
	defer pl.Close()

	go func() {
		conn, err := net.Dial("tcp", pl.Addr().String())
		if err != nil {
			t.Errorf("err: %v", err)
			return
		}
		defer conn.Close()

		conn.Write([]byte("PROXY TCP4 what 127.0.0.1 1000 2000\r\nping"))
	}()

	conn, err := pl.Accept()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer conn.Close()

	// Check the remote addr, should be the local addr
	addr := conn.RemoteAddr().(*net.TCPAddr)
	if addr.IP.String() != "127.0.0.1" {
		t.Fatalf("bad: %v", addr)
	}

	// Read should fail
	recv := make([]byte, 4)
	_, err = conn.Read(recv)
	if err == nil {
		t.Fatalf("err: %v", err)
	}
}
<|endoftext|>


<commit_before>package peerwriter

import (
	"bytes"
	"encoding/binary"
	"net"
	"time"

	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/internal/peer/peerprotocol"
	"github.com/cenkalti/rain/internal/piece"
)

const keepAlivePeriod = 2 * time.Minute

type PeerWriter struct {
	conn       net.Conn
	queueC     chan peerprotocol.Message
	writeQueue []peerprotocol.Message
	writeC     chan peerprotocol.Message
	log        logger.Logger
}

func New(conn net.Conn, l logger.Logger) *PeerWriter {
	return &PeerWriter{
		conn:       conn,
		queueC:     make(chan peerprotocol.Message),
		writeQueue: make([]peerprotocol.Message, 0),
		writeC:     make(chan peerprotocol.Message),
		log:        l,
	}
}

func (p *PeerWriter) SendMessage(msg peerprotocol.Message, stopC chan struct{}) {
	select {
	case p.queueC <- msg:
	case <-stopC:
	}
}

func (p *PeerWriter) SendPiece(msg peerprotocol.RequestMessage, pi *piece.Piece, stopC chan struct{}) {
	m := Piece{Piece: pi, Begin: msg.Begin, Length: msg.Length}
	select {
	case p.queueC <- m:
	case <-stopC:
	}
}

func (p *PeerWriter) Run(stopC chan struct{}) {
	go p.messageWriter(stopC)
	for {
		if len(p.writeQueue) == 0 {
			select {
			case msg := <-p.queueC:
				p.queueMessage(msg)
			case <-stopC:
				return
			}
		}
		msg := p.writeQueue[0]
		select {
		case msg = <-p.queueC:
			p.queueMessage(msg)
		case p.writeC <- msg:
			// TODO peer write queue array grows indefinitely. Try using linked list.
			p.writeQueue = p.writeQueue[1:]
		case <-stopC:
			return
		}
	}
}

func (p *PeerWriter) queueMessage(msg peerprotocol.Message) {
	// TODO cancel pending requests on choke
	p.writeQueue = append(p.writeQueue, msg)
}

func (p *PeerWriter) messageWriter(stopC chan struct{}) {
	keepAliveTicker := time.NewTicker(keepAlivePeriod)
	defer keepAliveTicker.Stop()
	for {
		select {
		case msg := <-p.writeC:
			p.log.Debugf("writing message of type: %q", msg.ID())
			payload, err := msg.MarshalBinary()
			if err != nil {
				p.log.Errorf("cannot marshal message [%v]: %s", msg.ID(), err.Error())
				p.conn.Close()
				return
			}
			buf := bytes.NewBuffer(make([]byte, 0, 4+1+len(payload)))
			var header = struct {
				Length uint32
				ID     peerprotocol.MessageID
			}{
				uint32(1 + len(payload)),
				msg.ID(),
			}
			_ = binary.Write(buf, binary.BigEndian, &header)
			buf.Write(payload)
			_, err = p.conn.Write(buf.Bytes())
			if err != nil {
				p.log.Errorf("cannot write message [%v]: %s", msg.ID(), err.Error())
				p.conn.Close()
				return
			}
		case <-keepAliveTicker.C:
			p.conn.Write([]byte{0, 0, 0, 0})
		case <-stopC:
			return
		}
	}
}
<commit_msg>close conn on keepalive error<commit_after>package peerwriter

import (
	"bytes"
	"encoding/binary"
	"net"
	"time"

	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/internal/peer/peerprotocol"
	"github.com/cenkalti/rain/internal/piece"
)

const keepAlivePeriod = 2 * time.Minute

type PeerWriter struct {
	conn       net.Conn
	queueC     chan peerprotocol.Message
	writeQueue []peerprotocol.Message
	writeC     chan peerprotocol.Message
	log        logger.Logger
}

func New(conn net.Conn, l logger.Logger) *PeerWriter {
	return &PeerWriter{
		conn:       conn,
		queueC:     make(chan peerprotocol.Message),
		writeQueue: make([]peerprotocol.Message, 0),
		writeC:     make(chan peerprotocol.Message),
		log:        l,
	}
}

func (p *PeerWriter) SendMessage(msg peerprotocol.Message, stopC chan struct{}) {
	select {
	case p.queueC <- msg:
	case <-stopC:
	}
}

func (p *PeerWriter) SendPiece(msg peerprotocol.RequestMessage, pi *piece.Piece, stopC chan struct{}) {
	m := Piece{Piece: pi, Begin: msg.Begin, Length: msg.Length}
	select {
	case p.queueC <- m:
	case <-stopC:
	}
}

func (p *PeerWriter) Run(stopC chan struct{}) {
	go p.messageWriter(stopC)
	for {
		if len(p.writeQueue) == 0 {
			select {
			case msg := <-p.queueC:
				p.queueMessage(msg)
			case <-stopC:
				return
			}
		}
		msg := p.writeQueue[0]
		select {
		case msg = <-p.queueC:
			p.queueMessage(msg)
		case p.writeC <- msg:
			// TODO peer write queue array grows indefinitely. Try using linked list.
			p.writeQueue = p.writeQueue[1:]
		case <-stopC:
			return
		}
	}
}

func (p *PeerWriter) queueMessage(msg peerprotocol.Message) {
	// TODO cancel pending requests on choke
	p.writeQueue = append(p.writeQueue, msg)
}

func (p *PeerWriter) messageWriter(stopC chan struct{}) {
	keepAliveTicker := time.NewTicker(keepAlivePeriod)
	defer keepAliveTicker.Stop()
	for {
		select {
		case msg := <-p.writeC:
			p.log.Debugf("writing message of type: %q", msg.ID())
			payload, err := msg.MarshalBinary()
			if err != nil {
				p.log.Errorf("cannot marshal message [%v]: %s", msg.ID(), err.Error())
				p.conn.Close()
				return
			}
			buf := bytes.NewBuffer(make([]byte, 0, 4+1+len(payload)))
			var header = struct {
				Length uint32
				ID     peerprotocol.MessageID
			}{
				uint32(1 + len(payload)),
				msg.ID(),
			}
			_ = binary.Write(buf, binary.BigEndian, &header)
			buf.Write(payload)
			_, err = p.conn.Write(buf.Bytes())
			if err != nil {
				p.log.Errorf("cannot write message [%v]: %s", msg.ID(), err.Error())
				p.conn.Close()
				return
			}
		case <-keepAliveTicker.C:
			_, err := p.conn.Write([]byte{0, 0, 0, 0})
			if err != nil {
				p.log.Errorf("cannot write keepalive message: %s", err.Error())
				p.conn.Close()
				return
			}
		case <-stopC:
			return
		}
	}
}
<|endoftext|>