<commit_before>package storage

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"
	"time"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/storage/backend"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	. "github.com/chrislusf/seaweedfs/weed/storage/types"
)

var ErrorNotFound = errors.New("not found")

// isFileUnchanged checks whether this needle to write is same as last one.
// It requires serialized access in the same volume.
func (v *Volume) isFileUnchanged(n *needle.Needle) bool {
	if v.Ttl.String() != "" {
		return false
	}

	nv, ok := v.nm.Get(n.Id)
	if ok && !nv.Offset.IsZero() && nv.Size != TombstoneFileSize {
		oldNeedle := new(needle.Needle)
		err := oldNeedle.ReadData(v.DataBackend, nv.Offset.ToAcutalOffset(), nv.Size, v.Version())
		if err != nil {
			glog.V(0).Infof("Failed to check updated file at offset %d size %d: %v", nv.Offset.ToAcutalOffset(), nv.Size, err)
			return false
		}
		if oldNeedle.Cookie == n.Cookie && oldNeedle.Checksum == n.Checksum && bytes.Equal(oldNeedle.Data, n.Data) {
			n.DataSize = oldNeedle.DataSize
			return true
		}
	}
	return false
}

// Destroy removes everything related to this volume
func (v *Volume) Destroy() (err error) {
	if v.isCompacting {
		err = fmt.Errorf("volume %d is compacting", v.Id)
		return
	}
	v.Close()
	os.Remove(v.FileName() + ".dat")
	os.Remove(v.FileName() + ".idx")
	os.Remove(v.FileName() + ".tier")
	os.Remove(v.FileName() + ".sdb")
	os.Remove(v.FileName() + ".cpd")
	os.Remove(v.FileName() + ".cpx")
	os.RemoveAll(v.FileName() + ".ldb")
	return
}

func (v *Volume) writeNeedle(n *needle.Needle) (offset uint64, size uint32, isUnchanged bool, err error) {
	glog.V(4).Infof("writing needle %s", needle.NewFileIdFromNeedle(v.Id, n).String())
	if v.readOnly {
		err = fmt.Errorf("%s is read-only", v.DataBackend.Name())
		return
	}
	v.dataFileAccessLock.Lock()
	defer v.dataFileAccessLock.Unlock()
	if v.isFileUnchanged(n) {
		size = n.DataSize
		isUnchanged = true
		return
	}

	if n.Ttl == needle.EMPTY_TTL && v.Ttl != needle.EMPTY_TTL {
		n.SetHasTtl()
		n.Ttl = v.Ttl
	}

	// check whether existing needle cookie matches
	nv, ok := v.nm.Get(n.Id)
	if ok {
		existingNeedle, _, _, existingNeedleReadErr := needle.ReadNeedleHeader(v.DataBackend, v.Version(), nv.Offset.ToAcutalOffset())
		if existingNeedleReadErr != nil {
			err = fmt.Errorf("reading existing needle: %v", existingNeedleReadErr)
			return
		}
		if existingNeedle.Cookie != n.Cookie {
			glog.V(0).Infof("write cookie mismatch: existing %x, new %x", existingNeedle.Cookie, n.Cookie)
			err = fmt.Errorf("mismatching cookie %x", n.Cookie)
			return
		}
	}

	// append to dat file
	n.AppendAtNs = uint64(time.Now().UnixNano())
	if offset, size, _, err = n.Append(v.DataBackend, v.Version()); err != nil {
		return
	}
	v.lastAppendAtNs = n.AppendAtNs

	// add to needle map
	if !ok || uint64(nv.Offset.ToAcutalOffset()) < offset {
		if err = v.nm.Put(n.Id, ToOffset(int64(offset)), n.Size); err != nil {
			glog.V(4).Infof("failed to save in needle map %d: %v", n.Id, err)
		}
	}
	if v.lastModifiedTsSeconds < n.LastModified {
		v.lastModifiedTsSeconds = n.LastModified
	}
	return
}

func (v *Volume) deleteNeedle(n *needle.Needle) (uint32, error) {
	glog.V(4).Infof("delete needle %s", needle.NewFileIdFromNeedle(v.Id, n).String())
	if v.readOnly {
		return 0, fmt.Errorf("%s is read-only", v.DataBackend.Name())
	}
	v.dataFileAccessLock.Lock()
	defer v.dataFileAccessLock.Unlock()
	nv, ok := v.nm.Get(n.Id)
	//fmt.Println("key", n.Id, "volume offset", nv.Offset, "data_size", n.Size, "cached size", nv.Size)
	if ok && nv.Size != TombstoneFileSize {
		size := nv.Size
		n.Data = nil
		n.AppendAtNs = uint64(time.Now().UnixNano())
		offset, _, _, err := n.Append(v.DataBackend, v.Version())
		if err != nil {
			return size, err
		}
		v.lastAppendAtNs = n.AppendAtNs
		if err = v.nm.Delete(n.Id, ToOffset(int64(offset))); err != nil {
			return size, err
		}
		return size, err
	}
	return 0, nil
}

// read fills in Needle content by looking up n.Id from NeedleMapper
func (v *Volume) readNeedle(n *needle.Needle) (int, error) {
	v.dataFileAccessLock.RLock()
	defer v.dataFileAccessLock.RUnlock()

	nv, ok := v.nm.Get(n.Id)
	if !ok || nv.Offset.IsZero() {
		return -1, ErrorNotFound
	}
	if nv.Size == TombstoneFileSize {
		return -1, errors.New("already deleted")
	}
	if nv.Size == 0 {
		return 0, nil
	}
	err := n.ReadData(v.DataBackend, nv.Offset.ToAcutalOffset(), nv.Size, v.Version())
	if err != nil {
		return 0, err
	}
	bytesRead := len(n.Data)
	if !n.HasTtl() {
		return bytesRead, nil
	}
	ttlMinutes := n.Ttl.Minutes()
	if ttlMinutes == 0 {
		return bytesRead, nil
	}
	if !n.HasLastModifiedDate() {
		return bytesRead, nil
	}
	if uint64(time.Now().Unix()) < n.LastModified+uint64(ttlMinutes*60) {
		return bytesRead, nil
	}
	return -1, ErrorNotFound
}

type VolumeFileScanner interface {
	VisitSuperBlock(SuperBlock) error
	ReadNeedleBody() bool
	VisitNeedle(n *needle.Needle, offset int64, needleHeader, needleBody []byte) error
}

func ScanVolumeFile(dirname string, collection string, id needle.VolumeId,
	needleMapKind NeedleMapType,
	volumeFileScanner VolumeFileScanner) (err error) {
	var v *Volume
	if v, err = loadVolumeWithoutIndex(dirname, collection, id, needleMapKind); err != nil {
		return fmt.Errorf("failed to load volume %d: %v", id, err)
	}
	if err = volumeFileScanner.VisitSuperBlock(v.SuperBlock); err != nil {
		return fmt.Errorf("failed to process volume %d super block: %v", id, err)
	}
	defer v.Close()

	version := v.Version()

	offset := int64(v.SuperBlock.BlockSize())

	return ScanVolumeFileFrom(version, v.DataBackend, offset, volumeFileScanner)
}

func ScanVolumeFileFrom(version needle.Version, datBackend backend.BackendStorageFile, offset int64, volumeFileScanner VolumeFileScanner) (err error) {
	n, nh, rest, e := needle.ReadNeedleHeader(datBackend, version, offset)
	if e != nil {
		if e == io.EOF {
			return nil
		}
		return fmt.Errorf("cannot read %s at offset %d: %v", datBackend.Name(), offset, e)
	}
	for n != nil {
		var needleBody []byte
		if volumeFileScanner.ReadNeedleBody() {
			if needleBody, err = n.ReadNeedleBody(datBackend, version, offset+NeedleHeaderSize, rest); err != nil {
				glog.V(0).Infof("cannot read needle body: %v", err)
				//err = fmt.Errorf("cannot read needle body: %v", err)
				//return
			}
		}
		err := volumeFileScanner.VisitNeedle(n, offset, nh, needleBody)
		if err == io.EOF {
			return nil
		}
		if err != nil {
			glog.V(0).Infof("visit needle error: %v", err)
			return fmt.Errorf("visit needle error: %v", err)
		}
		offset += NeedleHeaderSize + rest
		glog.V(4).Infof("==> new entry offset %d", offset)
		if n, nh, rest, err = needle.ReadNeedleHeader(datBackend, version, offset); err != nil {
			if err == io.EOF {
				return nil
			}
			return fmt.Errorf("cannot read needle header at offset %d: %v", offset, err)
		}
		glog.V(4).Infof("new entry needle size:%d rest:%d", n.Size, rest)
	}
	return nil
}
<commit_msg>remove duplicated checking<commit_after>package storage

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"
	"time"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/storage/backend"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	. "github.com/chrislusf/seaweedfs/weed/storage/types"
)

var ErrorNotFound = errors.New("not found")

// isFileUnchanged checks whether this needle to write is same as last one.
// It requires serialized access in the same volume.
func (v *Volume) isFileUnchanged(n *needle.Needle) bool {
	if v.Ttl.String() != "" {
		return false
	}

	nv, ok := v.nm.Get(n.Id)
	if ok && !nv.Offset.IsZero() && nv.Size != TombstoneFileSize {
		oldNeedle := new(needle.Needle)
		err := oldNeedle.ReadData(v.DataBackend, nv.Offset.ToAcutalOffset(), nv.Size, v.Version())
		if err != nil {
			glog.V(0).Infof("Failed to check updated file at offset %d size %d: %v", nv.Offset.ToAcutalOffset(), nv.Size, err)
			return false
		}
		if oldNeedle.Cookie == n.Cookie && oldNeedle.Checksum == n.Checksum && bytes.Equal(oldNeedle.Data, n.Data) {
			n.DataSize = oldNeedle.DataSize
			return true
		}
	}
	return false
}

// Destroy removes everything related to this volume
func (v *Volume) Destroy() (err error) {
	if v.isCompacting {
		err = fmt.Errorf("volume %d is compacting", v.Id)
		return
	}
	v.Close()
	os.Remove(v.FileName() + ".dat")
	os.Remove(v.FileName() + ".idx")
	os.Remove(v.FileName() + ".tier")
	os.Remove(v.FileName() + ".sdb")
	os.Remove(v.FileName() + ".cpd")
	os.Remove(v.FileName() + ".cpx")
	os.RemoveAll(v.FileName() + ".ldb")
	return
}

func (v *Volume) writeNeedle(n *needle.Needle) (offset uint64, size uint32, isUnchanged bool, err error) {
	glog.V(4).Infof("writing needle %s", needle.NewFileIdFromNeedle(v.Id, n).String())
	v.dataFileAccessLock.Lock()
	defer v.dataFileAccessLock.Unlock()
	if v.isFileUnchanged(n) {
		size = n.DataSize
		isUnchanged = true
		return
	}

	if n.Ttl == needle.EMPTY_TTL && v.Ttl != needle.EMPTY_TTL {
		n.SetHasTtl()
		n.Ttl = v.Ttl
	}

	// check whether existing needle cookie matches
	nv, ok := v.nm.Get(n.Id)
	if ok {
		existingNeedle, _, _, existingNeedleReadErr := needle.ReadNeedleHeader(v.DataBackend, v.Version(), nv.Offset.ToAcutalOffset())
		if existingNeedleReadErr != nil {
			err = fmt.Errorf("reading existing needle: %v", existingNeedleReadErr)
			return
		}
		if existingNeedle.Cookie != n.Cookie {
			glog.V(0).Infof("write cookie mismatch: existing %x, new %x", existingNeedle.Cookie, n.Cookie)
			err = fmt.Errorf("mismatching cookie %x", n.Cookie)
			return
		}
	}

	// append to dat file
	n.AppendAtNs = uint64(time.Now().UnixNano())
	if offset, size, _, err = n.Append(v.DataBackend, v.Version()); err != nil {
		return
	}
	v.lastAppendAtNs = n.AppendAtNs

	// add to needle map
	if !ok || uint64(nv.Offset.ToAcutalOffset()) < offset {
		if err = v.nm.Put(n.Id, ToOffset(int64(offset)), n.Size); err != nil {
			glog.V(4).Infof("failed to save in needle map %d: %v", n.Id, err)
		}
	}
	if v.lastModifiedTsSeconds < n.LastModified {
		v.lastModifiedTsSeconds = n.LastModified
	}
	return
}

func (v *Volume) deleteNeedle(n *needle.Needle) (uint32, error) {
	glog.V(4).Infof("delete needle %s", needle.NewFileIdFromNeedle(v.Id, n).String())
	v.dataFileAccessLock.Lock()
	defer v.dataFileAccessLock.Unlock()
	nv, ok := v.nm.Get(n.Id)
	//fmt.Println("key", n.Id, "volume offset", nv.Offset, "data_size", n.Size, "cached size", nv.Size)
	if ok && nv.Size != TombstoneFileSize {
		size := nv.Size
		n.Data = nil
		n.AppendAtNs = uint64(time.Now().UnixNano())
		offset, _, _, err := n.Append(v.DataBackend, v.Version())
		if err != nil {
			return size, err
		}
		v.lastAppendAtNs = n.AppendAtNs
		if err = v.nm.Delete(n.Id, ToOffset(int64(offset))); err != nil {
			return size, err
		}
		return size, err
	}
	return 0, nil
}

// read fills in Needle content by looking up n.Id from NeedleMapper
func (v *Volume) readNeedle(n *needle.Needle) (int, error) {
	v.dataFileAccessLock.RLock()
	defer v.dataFileAccessLock.RUnlock()

	nv, ok := v.nm.Get(n.Id)
	if !ok || nv.Offset.IsZero() {
		return -1, ErrorNotFound
	}
	if nv.Size == TombstoneFileSize {
		return -1, errors.New("already deleted")
	}
	if nv.Size == 0 {
		return 0, nil
	}
	err := n.ReadData(v.DataBackend, nv.Offset.ToAcutalOffset(), nv.Size, v.Version())
	if err != nil {
		return 0, err
	}
	bytesRead := len(n.Data)
	if !n.HasTtl() {
		return bytesRead, nil
	}
	ttlMinutes := n.Ttl.Minutes()
	if ttlMinutes == 0 {
		return bytesRead, nil
	}
	if !n.HasLastModifiedDate() {
		return bytesRead, nil
	}
	if uint64(time.Now().Unix()) < n.LastModified+uint64(ttlMinutes*60) {
		return bytesRead, nil
	}
	return -1, ErrorNotFound
}

type VolumeFileScanner interface {
	VisitSuperBlock(SuperBlock) error
	ReadNeedleBody() bool
	VisitNeedle(n *needle.Needle, offset int64, needleHeader, needleBody []byte) error
}

func ScanVolumeFile(dirname string, collection string, id needle.VolumeId,
	needleMapKind NeedleMapType,
	volumeFileScanner VolumeFileScanner) (err error) {
	var v *Volume
	if v, err = loadVolumeWithoutIndex(dirname, collection, id, needleMapKind); err != nil {
		return fmt.Errorf("failed to load volume %d: %v", id, err)
	}
	if err = volumeFileScanner.VisitSuperBlock(v.SuperBlock); err != nil {
		return fmt.Errorf("failed to process volume %d super block: %v", id, err)
	}
	defer v.Close()

	version := v.Version()

	offset := int64(v.SuperBlock.BlockSize())

	return ScanVolumeFileFrom(version, v.DataBackend, offset, volumeFileScanner)
}

func ScanVolumeFileFrom(version needle.Version, datBackend backend.BackendStorageFile, offset int64, volumeFileScanner VolumeFileScanner) (err error) {
	n, nh, rest, e := needle.ReadNeedleHeader(datBackend, version, offset)
	if e != nil {
		if e == io.EOF {
			return nil
		}
		return fmt.Errorf("cannot read %s at offset %d: %v", datBackend.Name(), offset, e)
	}
	for n != nil {
		var needleBody []byte
		if volumeFileScanner.ReadNeedleBody() {
			if needleBody, err = n.ReadNeedleBody(datBackend, version, offset+NeedleHeaderSize, rest); err != nil {
				glog.V(0).Infof("cannot read needle body: %v", err)
				//err = fmt.Errorf("cannot read needle body: %v", err)
				//return
			}
		}
		err := volumeFileScanner.VisitNeedle(n, offset, nh, needleBody)
		if err == io.EOF {
			return nil
		}
		if err != nil {
			glog.V(0).Infof("visit needle error: %v", err)
			return fmt.Errorf("visit needle error: %v", err)
		}
		offset += NeedleHeaderSize + rest
		glog.V(4).Infof("==> new entry offset %d", offset)
		if n, nh, rest, err = needle.ReadNeedleHeader(datBackend, version, offset); err != nil {
			if err == io.EOF {
				return nil
			}
			return fmt.Errorf("cannot read needle header at offset %d: %v", offset, err)
		}
		glog.V(4).Infof("new entry needle size:%d rest:%d", n.Size, rest)
	}
	return nil
}
<|endoftext|>


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

package webpagereplay

import (
	"crypto"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io"
	"net"
	"sync"
	"time"
)

// Returns a TLS configuration that serves a recorded server leaf cert signed by
// root CA.
func ReplayTLSConfig(root tls.Certificate, a *Archive) (*tls.Config, error) {
	root_cert, err := getRootCert(root)
	if err != nil {
		return nil, fmt.Errorf("bad local cert: %v", err)
	}
	tp := &tlsProxy{&root, root_cert, a, nil, sync.Mutex{}, make(map[string][]byte)}
	return &tls.Config{
		GetConfigForClient: tp.getReplayConfigForClient,
	}, nil
}

// Returns a TLS configuration that serves a server leaf cert fetched over the
// network on demand.
func RecordTLSConfig(root tls.Certificate, w *WritableArchive) (*tls.Config, error) {
	root_cert, err := getRootCert(root)
	if err != nil {
		return nil, fmt.Errorf("bad local cert: %v", err)
	}
	tp := &tlsProxy{&root, root_cert, nil, w, sync.Mutex{}, nil}
	return &tls.Config{
		GetConfigForClient: tp.getRecordConfigForClient,
	}, nil
}

func getRootCert(root tls.Certificate) (*x509.Certificate, error) {
	root_cert, err := x509.ParseCertificate(root.Certificate[0])
	if err != nil {
		return nil, err
	}
	root_cert.IsCA = true
	root_cert.BasicConstraintsValid = true
	return root_cert, nil
}

// Mints a dummy server cert when the real one is not recorded.
func MintDummyCertificate(serverName string, rootCert *x509.Certificate, rootKey crypto.PrivateKey) ([]byte, string, error) {
	template := rootCert
	if ip := net.ParseIP(serverName); ip != nil {
		template.IPAddresses = []net.IP{ip}
	} else {
		template.DNSNames = []string{serverName}
	}
	var buf [20]byte
	if _, err := io.ReadFull(rand.Reader, buf[:]); err != nil {
		return nil, "", fmt.Errorf("create cert failed: %v", err)
	}
	template.SerialNumber.SetBytes(buf[:])
	template.Issuer = template.Subject
	derBytes, err := x509.CreateCertificate(rand.Reader, template, template, template.PublicKey, rootKey)
	if err != nil {
		return nil, "", fmt.Errorf("create cert failed: %v", err)
	}
	return derBytes, "", err
}

// Returns DER encoded server cert.
func MintServerCert(serverName string, rootCert *x509.Certificate, rootKey crypto.PrivateKey) ([]byte, string, error) {
	dialer := &net.Dialer{
		Timeout:   30 * time.Second,
		KeepAlive: 30 * time.Second,
		DualStack: true,
	}
	conn, err := tls.DialWithDialer(dialer, "tcp", fmt.Sprintf("%s:443", serverName), &tls.Config{
		NextProtos: []string{"h2", "http/1.1"},
	})
	if err != nil {
		return nil, "", fmt.Errorf("Couldn't reach host %s: %v", serverName, err)
	}
	defer conn.Close()
	conn.Handshake()
	template := conn.ConnectionState().PeerCertificates[0]

	template.Subject.CommonName = serverName
	template.NotBefore = time.Now()
	// Certs cannot be valid for longer than 39 mths.
	template.NotAfter = template.NotBefore.Add(39 * 30 * 24 * time.Hour)
	template.SignatureAlgorithm = rootCert.SignatureAlgorithm
	template.PublicKey = rootCert.PublicKey
	var buf [20]byte
	if _, err := io.ReadFull(rand.Reader, buf[:]); err != nil {
		return nil, "", err
	}
	template.SerialNumber.SetBytes(buf[:])
	template.Issuer = rootCert.Subject
	template.KeyUsage = x509.KeyUsageCertSign | x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature | x509.KeyUsageCRLSign
	template.ExtKeyUsage = []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth, x509.ExtKeyUsageServerAuth}

	negotiatedProtocol := conn.ConnectionState().NegotiatedProtocol
	derBytes, err := x509.CreateCertificate(rand.Reader, template, rootCert, template.PublicKey, rootKey)
	return derBytes, negotiatedProtocol, err
}

type tlsProxy struct {
	root             *tls.Certificate
	root_cert        *x509.Certificate
	archive          *Archive
	writable_archive *WritableArchive
	mu               sync.Mutex
	dummy_certs_map  map[string][]byte
}

// TODO: For now, this just returns a self-signed cert using the given ServerName.
// In the future, for better HTTP/2 support, we may want to record host equivalence
// classes in the archive, where an equivalence class contains all hosts that can be
// served by the same IP. We can then run a DNS proxy that maps all hostnames in the
// same equivalence class to the same local port, which models the possibility that
// every equivalence class of hostnames can be served over the same HTTP/2 connection.
func (tp *tlsProxy) getReplayConfigForClient(clientHello *tls.ClientHelloInfo) (*tls.Config, error) {
	h := clientHello.ServerName
	if h == "" {
		return &tls.Config{
			Certificates: []tls.Certificate{*tp.root},
		}, nil
	}

	derBytes, negotiatedProtocol, err := tp.archive.FindHostTlsConfig(h)
	tp.mu.Lock()
	defer tp.mu.Unlock()
	if err != nil || derBytes == nil {
		if _, ok := tp.dummy_certs_map[h]; !ok {
			derBytes, negotiatedProtocol, err = MintDummyCertificate(h, tp.root_cert, tp.root.PrivateKey)
			if err != nil {
				return nil, err
			}
			tp.dummy_certs_map[h] = derBytes
		}
		derBytes = tp.dummy_certs_map[h]
	}
	return &tls.Config{
		Certificates: []tls.Certificate{
			tls.Certificate{
				Certificate: [][]byte{derBytes},
				PrivateKey:  tp.root.PrivateKey,
			}},
		NextProtos: buildNextProtos(negotiatedProtocol),
	}, nil
}

func buildNextProtos(negotiatedProtocol string) []string {
	if negotiatedProtocol == "h2" {
		return []string{"h2", "http/1.1"}
	}
	return []string{"http/1.1"}
}

func (tp *tlsProxy) getRecordConfigForClient(clientHello *tls.ClientHelloInfo) (*tls.Config, error) {
	h := clientHello.ServerName
	if h == "" {
		return &tls.Config{
			Certificates: []tls.Certificate{*tp.root},
		}, nil
	}
	derBytes, negotiatedProtocol, err := tp.writable_archive.Archive.FindHostTlsConfig(h)
	if err == nil && derBytes != nil {
		return &tls.Config{
			Certificates: []tls.Certificate{
				tls.Certificate{
					Certificate: [][]byte{derBytes},
					PrivateKey:  tp.root.PrivateKey,
				}},
			NextProtos: buildNextProtos(negotiatedProtocol),
		}, nil
	}

	derBytes, negotiatedProtocol, err = MintServerCert(h, tp.root_cert, tp.root.PrivateKey)
	if err != nil {
		return nil, fmt.Errorf("create cert failed: %v", err)
	}

	tp.writable_archive.RecordTlsConfig(h, derBytes, negotiatedProtocol)

	return &tls.Config{
		Certificates: []tls.Certificate{
			tls.Certificate{
				Certificate: [][]byte{derBytes},
				PrivateKey:  tp.root.PrivateKey}},
		NextProtos: buildNextProtos(negotiatedProtocol),
	}, nil
}
<commit_msg>Changed mint certificate expiration date to no more than 12 months<commit_after>// Copyright 2017 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package webpagereplay

import (
	"crypto"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io"
	"net"
	"sync"
	"time"
)

// Returns a TLS configuration that serves a recorded server leaf cert signed by
// root CA.
func ReplayTLSConfig(root tls.Certificate, a *Archive) (*tls.Config, error) {
	root_cert, err := getRootCert(root)
	if err != nil {
		return nil, fmt.Errorf("bad local cert: %v", err)
	}
	tp := &tlsProxy{&root, root_cert, a, nil, sync.Mutex{}, make(map[string][]byte)}
	return &tls.Config{
		GetConfigForClient: tp.getReplayConfigForClient,
	}, nil
}

// Returns a TLS configuration that serves a server leaf cert fetched over the
// network on demand.
func RecordTLSConfig(root tls.Certificate, w *WritableArchive) (*tls.Config, error) {
	root_cert, err := getRootCert(root)
	if err != nil {
		return nil, fmt.Errorf("bad local cert: %v", err)
	}
	tp := &tlsProxy{&root, root_cert, nil, w, sync.Mutex{}, nil}
	return &tls.Config{
		GetConfigForClient: tp.getRecordConfigForClient,
	}, nil
}

func getRootCert(root tls.Certificate) (*x509.Certificate, error) {
	root_cert, err := x509.ParseCertificate(root.Certificate[0])
	if err != nil {
		return nil, err
	}
	root_cert.IsCA = true
	root_cert.BasicConstraintsValid = true
	return root_cert, nil
}

// Mints a dummy server cert when the real one is not recorded.
func MintDummyCertificate(serverName string, rootCert *x509.Certificate, rootKey crypto.PrivateKey) ([]byte, string, error) {
	template := rootCert
	if ip := net.ParseIP(serverName); ip != nil {
		template.IPAddresses = []net.IP{ip}
	} else {
		template.DNSNames = []string{serverName}
	}
	var buf [20]byte
	if _, err := io.ReadFull(rand.Reader, buf[:]); err != nil {
		return nil, "", fmt.Errorf("create cert failed: %v", err)
	}
	template.SerialNumber.SetBytes(buf[:])
	template.Issuer = template.Subject
	derBytes, err := x509.CreateCertificate(rand.Reader, template, template, template.PublicKey, rootKey)
	if err != nil {
		return nil, "", fmt.Errorf("create cert failed: %v", err)
	}
	return derBytes, "", err
}

// Returns DER encoded server cert.
func MintServerCert(serverName string, rootCert *x509.Certificate, rootKey crypto.PrivateKey) ([]byte, string, error) {
	dialer := &net.Dialer{
		Timeout:   30 * time.Second,
		KeepAlive: 30 * time.Second,
		DualStack: true,
	}
	conn, err := tls.DialWithDialer(dialer, "tcp", fmt.Sprintf("%s:443", serverName), &tls.Config{
		NextProtos: []string{"h2", "http/1.1"},
	})
	if err != nil {
		return nil, "", fmt.Errorf("Couldn't reach host %s: %v", serverName, err)
	}
	defer conn.Close()
	conn.Handshake()
	template := conn.ConnectionState().PeerCertificates[0]

	template.Subject.CommonName = serverName
	template.NotBefore = time.Now()
	// Certs cannot be valid for longer than 12 mths.
	template.NotAfter = template.NotBefore.Add(12 * 30 * 24 * time.Hour)
	template.SignatureAlgorithm = rootCert.SignatureAlgorithm
	template.PublicKey = rootCert.PublicKey
	var buf [20]byte
	if _, err := io.ReadFull(rand.Reader, buf[:]); err != nil {
		return nil, "", err
	}
	template.SerialNumber.SetBytes(buf[:])
	template.Issuer = rootCert.Subject
	template.KeyUsage = x509.KeyUsageCertSign | x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature | x509.KeyUsageCRLSign
	template.ExtKeyUsage = []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth, x509.ExtKeyUsageServerAuth}

	negotiatedProtocol := conn.ConnectionState().NegotiatedProtocol
	derBytes, err := x509.CreateCertificate(rand.Reader, template, rootCert, template.PublicKey, rootKey)
	return derBytes, negotiatedProtocol, err
}

type tlsProxy struct {
	root             *tls.Certificate
	root_cert        *x509.Certificate
	archive          *Archive
	writable_archive *WritableArchive
	mu               sync.Mutex
	dummy_certs_map  map[string][]byte
}

// TODO: For now, this just returns a self-signed cert using the given ServerName.
// In the future, for better HTTP/2 support, we may want to record host equivalence
// classes in the archive, where an equivalence class contains all hosts that can be
// served by the same IP. We can then run a DNS proxy that maps all hostnames in the
// same equivalence class to the same local port, which models the possibility that
// every equivalence class of hostnames can be served over the same HTTP/2 connection.
func (tp *tlsProxy) getReplayConfigForClient(clientHello *tls.ClientHelloInfo) (*tls.Config, error) {
	h := clientHello.ServerName
	if h == "" {
		return &tls.Config{
			Certificates: []tls.Certificate{*tp.root},
		}, nil
	}

	derBytes, negotiatedProtocol, err := tp.archive.FindHostTlsConfig(h)
	tp.mu.Lock()
	defer tp.mu.Unlock()
	if err != nil || derBytes == nil {
		if _, ok := tp.dummy_certs_map[h]; !ok {
			derBytes, negotiatedProtocol, err = MintDummyCertificate(h, tp.root_cert, tp.root.PrivateKey)
			if err != nil {
				return nil, err
			}
			tp.dummy_certs_map[h] = derBytes
		}
		derBytes = tp.dummy_certs_map[h]
	}
	return &tls.Config{
		Certificates: []tls.Certificate{
			tls.Certificate{
				Certificate: [][]byte{derBytes},
				PrivateKey:  tp.root.PrivateKey,
			}},
		NextProtos: buildNextProtos(negotiatedProtocol),
	}, nil
}

func buildNextProtos(negotiatedProtocol string) []string {
	if negotiatedProtocol == "h2" {
		return []string{"h2", "http/1.1"}
	}
	return []string{"http/1.1"}
}

func (tp *tlsProxy) getRecordConfigForClient(clientHello *tls.ClientHelloInfo) (*tls.Config, error) {
	h := clientHello.ServerName
	if h == "" {
		return &tls.Config{
			Certificates: []tls.Certificate{*tp.root},
		}, nil
	}
	derBytes, negotiatedProtocol, err := tp.writable_archive.Archive.FindHostTlsConfig(h)
	if err == nil && derBytes != nil {
		return &tls.Config{
			Certificates: []tls.Certificate{
				tls.Certificate{
					Certificate: [][]byte{derBytes},
					PrivateKey:  tp.root.PrivateKey,
				}},
			NextProtos: buildNextProtos(negotiatedProtocol),
		}, nil
	}

	derBytes, negotiatedProtocol, err = MintServerCert(h, tp.root_cert, tp.root.PrivateKey)
	if err != nil {
		return nil, fmt.Errorf("create cert failed: %v", err)
	}

	tp.writable_archive.RecordTlsConfig(h, derBytes, negotiatedProtocol)

	return &tls.Config{
		Certificates: []tls.Certificate{
			tls.Certificate{
				Certificate: [][]byte{derBytes},
				PrivateKey:  tp.root.PrivateKey}},
		NextProtos: buildNextProtos(negotiatedProtocol),
	}, nil
}
<|endoftext|>


<commit_before>package abstract_sql

import (
	"context"
	"database/sql"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/filer"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/util"
)

type AbstractSqlStore struct {
	DB                      *sql.DB
	SqlInsert               string
	SqlUpdate               string
	SqlFind                 string
	SqlDelete               string
	SqlDeleteFolderChildren string
	SqlListExclusive        string
	SqlListInclusive        string
}

type TxOrDB interface {
	ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error)
	QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row
	QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error)
}

func (store *AbstractSqlStore) BeginTransaction(ctx context.Context) (context.Context, error) {
	tx, err := store.DB.BeginTx(ctx, &sql.TxOptions{
		Isolation: sql.LevelReadCommitted,
		ReadOnly:  false,
	})
	if err != nil {
		return ctx, err
	}

	return context.WithValue(ctx, "tx", tx), nil
}
func (store *AbstractSqlStore) CommitTransaction(ctx context.Context) error {
	if tx, ok := ctx.Value("tx").(*sql.Tx); ok {
		return tx.Commit()
	}
	return nil
}
func (store *AbstractSqlStore) RollbackTransaction(ctx context.Context) error {
	if tx, ok := ctx.Value("tx").(*sql.Tx); ok {
		return tx.Rollback()
	}
	return nil
}

func (store *AbstractSqlStore) getTxOrDB(ctx context.Context) TxOrDB {
	if tx, ok := ctx.Value("tx").(*sql.Tx); ok {
		return tx
	}
	return store.DB
}

func (store *AbstractSqlStore) InsertEntry(ctx context.Context, entry *filer.Entry) (err error) {

	dir, name := entry.FullPath.DirAndName()
	meta, err := entry.EncodeAttributesAndChunks()
	if err != nil {
		return fmt.Errorf("encode %s: %s", entry.FullPath, err)
	}

	res, err := store.getTxOrDB(ctx).ExecContext(ctx, store.SqlInsert, util.HashStringToLong(dir), name, dir, meta)
	if err != nil {
		return fmt.Errorf("insert %s: %s", entry.FullPath, err)
	}

	affectedRows, err := res.RowsAffected()
	if err == nil && affectedRows > 0 {
		return nil
	}

	// now the insert failed possibly due to duplication constraints
	glog.V(1).Infof("insert %s falls back to update: %s", entry.FullPath, err)

	res, err = store.getTxOrDB(ctx).ExecContext(ctx, store.SqlUpdate, meta, util.HashStringToLong(dir), name, dir)
	if err != nil {
		return fmt.Errorf("upsert %s: %s", entry.FullPath, err)
	}

	_, err = res.RowsAffected()
	if err != nil {
		return fmt.Errorf("upsert %s but no rows affected: %s", entry.FullPath, err)
	}
	return nil

}

func (store *AbstractSqlStore) UpdateEntry(ctx context.Context, entry *filer.Entry) (err error) {

	dir, name := entry.FullPath.DirAndName()
	meta, err := entry.EncodeAttributesAndChunks()
	if err != nil {
		return fmt.Errorf("encode %s: %s", entry.FullPath, err)
	}

	res, err := store.getTxOrDB(ctx).ExecContext(ctx, store.SqlUpdate, meta, util.HashStringToLong(dir), name, dir)
	if err != nil {
		return fmt.Errorf("update %s: %s", entry.FullPath, err)
	}

	_, err = res.RowsAffected()
	if err != nil {
		return fmt.Errorf("update %s but no rows affected: %s", entry.FullPath, err)
	}
	return nil
}

func (store *AbstractSqlStore) FindEntry(ctx context.Context, fullpath util.FullPath) (*filer.Entry, error) {

	dir, name := fullpath.DirAndName()
	row := store.getTxOrDB(ctx).QueryRowContext(ctx, store.SqlFind, util.HashStringToLong(dir), name, dir)
	var data []byte
	if err := row.Scan(&data); err != nil {
		return nil, filer_pb.ErrNotFound
	}

	entry := &filer.Entry{
		FullPath: fullpath,
	}
	if err := entry.DecodeAttributesAndChunks(data); err != nil {
		return entry, fmt.Errorf("decode %s : %v", entry.FullPath, err)
	}

	return entry, nil
}

func (store *AbstractSqlStore) DeleteEntry(ctx context.Context, fullpath util.FullPath) error {

	dir, name := fullpath.DirAndName()

	res, err := store.getTxOrDB(ctx).ExecContext(ctx, store.SqlDelete, util.HashStringToLong(dir), name, dir)
	if err != nil {
		return fmt.Errorf("delete %s: %s", fullpath, err)
	}

	_, err = res.RowsAffected()
	if err != nil {
		return fmt.Errorf("delete %s but no rows affected: %s", fullpath, err)
	}

	return nil
}

func (store *AbstractSqlStore) DeleteFolderChildren(ctx context.Context, fullpath util.FullPath) error {

	res, err := store.getTxOrDB(ctx).ExecContext(ctx, store.SqlDeleteFolderChildren, util.HashStringToLong(string(fullpath)), fullpath)
	if err != nil {
		return fmt.Errorf("deleteFolderChildren %s: %s", fullpath, err)
	}

	_, err = res.RowsAffected()
	if err != nil {
		return fmt.Errorf("deleteFolderChildren %s but no rows affected: %s", fullpath, err)
	}

	return nil
}

func (store *AbstractSqlStore) ListDirectoryPrefixedEntries(ctx context.Context, fullpath util.FullPath, startFileName string, inclusive bool, limit int, prefix string) (entries []*filer.Entry, err error) {
	sqlText := store.SqlListExclusive
	if inclusive {
		sqlText = store.SqlListInclusive
	}

	rows, err := store.getTxOrDB(ctx).QueryContext(ctx, sqlText, util.HashStringToLong(string(fullpath)), startFileName, string(fullpath), prefix, limit)
	if err != nil {
		return nil, fmt.Errorf("list %s : %v", fullpath, err)
	}
	defer rows.Close()

	for rows.Next() {
		var name string
		var data []byte
		if err = rows.Scan(&name, &data); err != nil {
			glog.V(0).Infof("scan %s : %v", fullpath, err)
			return nil, fmt.Errorf("scan %s: %v", fullpath, err)
		}

		entry := &filer.Entry{
			FullPath: util.NewFullPath(string(fullpath), name),
		}
		if err = entry.DecodeAttributesAndChunks(data); err != nil {
			glog.V(0).Infof("scan decode %s : %v", entry.FullPath, err)
			return nil, fmt.Errorf("scan decode %s : %v", entry.FullPath, err)
		}

		entries = append(entries, entry)
	}

	return entries, nil
}

func (store *AbstractSqlStore) ListDirectoryEntries(ctx context.Context, fullpath util.FullPath, startFileName string, inclusive bool, limit int) (entries []*filer.Entry, err error) {
	return store.ListDirectoryPrefixedEntries(ctx, fullpath, startFileName, inclusive, limit, "")
}

func (store *AbstractSqlStore) Shutdown() {
	store.DB.Close()
}
<commit_msg>mysql or postgres: log find error<commit_after>package abstract_sql

import (
	"context"
	"database/sql"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/filer"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/util"
)

type AbstractSqlStore struct {
	DB                      *sql.DB
	SqlInsert               string
	SqlUpdate               string
	SqlFind                 string
	SqlDelete               string
	SqlDeleteFolderChildren string
	SqlListExclusive        string
	SqlListInclusive        string
}

type TxOrDB interface {
	ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error)
	QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row
	QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error)
}

func (store *AbstractSqlStore) BeginTransaction(ctx context.Context) (context.Context, error) {
	tx, err := store.DB.BeginTx(ctx, &sql.TxOptions{
		Isolation: sql.LevelReadCommitted,
		ReadOnly:  false,
	})
	if err != nil {
		return ctx, err
	}

	return context.WithValue(ctx, "tx", tx), nil
}
func (store *AbstractSqlStore) CommitTransaction(ctx context.Context) error {
	if tx, ok := ctx.Value("tx").(*sql.Tx); ok {
		return tx.Commit()
	}
	return nil
}
func (store *AbstractSqlStore) RollbackTransaction(ctx context.Context) error {
	if tx, ok := ctx.Value("tx").(*sql.Tx); ok {
		return tx.Rollback()
	}
	return nil
}

func (store *AbstractSqlStore) getTxOrDB(ctx context.Context) TxOrDB {
	if tx, ok := ctx.Value("tx").(*sql.Tx); ok {
		return tx
	}
	return store.DB
}

func (store *AbstractSqlStore) InsertEntry(ctx context.Context, entry *filer.Entry) (err error) {

	dir, name := entry.FullPath.DirAndName()
	meta, err := entry.EncodeAttributesAndChunks()
	if err != nil {
		return fmt.Errorf("encode %s: %s", entry.FullPath, err)
	}

	res, err := store.getTxOrDB(ctx).ExecContext(ctx, store.SqlInsert, util.HashStringToLong(dir), name, dir, meta)
	if err != nil {
		return fmt.Errorf("insert %s: %s", entry.FullPath, err)
	}

	affectedRows, err := res.RowsAffected()
	if err == nil && affectedRows > 0 {
		return nil
	}

	// now the insert failed possibly due to duplication constraints
	glog.V(1).Infof("insert %s falls back to update: %s", entry.FullPath, err)

	res, err = store.getTxOrDB(ctx).ExecContext(ctx, store.SqlUpdate, meta, util.HashStringToLong(dir), name, dir)
	if err != nil {
		return fmt.Errorf("upsert %s: %s", entry.FullPath, err)
	}

	_, err = res.RowsAffected()
	if err != nil {
		return fmt.Errorf("upsert %s but no rows affected: %s", entry.FullPath, err)
	}
	return nil

}

func (store *AbstractSqlStore) UpdateEntry(ctx context.Context, entry *filer.Entry) (err error) {

	dir, name := entry.FullPath.DirAndName()
	meta, err := entry.EncodeAttributesAndChunks()
	if err != nil {
		return fmt.Errorf("encode %s: %s", entry.FullPath, err)
	}

	res, err := store.getTxOrDB(ctx).ExecContext(ctx, store.SqlUpdate, meta, util.HashStringToLong(dir), name, dir)
	if err != nil {
		return fmt.Errorf("update %s: %s", entry.FullPath, err)
	}

	_, err = res.RowsAffected()
	if err != nil {
		return fmt.Errorf("update %s but no rows affected: %s", entry.FullPath, err)
	}
	return nil
}

func (store *AbstractSqlStore) FindEntry(ctx context.Context, fullpath util.FullPath) (*filer.Entry, error) {

	dir, name := fullpath.DirAndName()
	row := store.getTxOrDB(ctx).QueryRowContext(ctx, store.SqlFind, util.HashStringToLong(dir), name, dir)
	var data []byte
	if err := row.Scan(&data); err != nil {
		glog.Errorf("find %s: %v", fullpath, err)
		return nil, filer_pb.ErrNotFound
	}

	entry := &filer.Entry{
		FullPath: fullpath,
	}
	if err := entry.DecodeAttributesAndChunks(data); err != nil {
		return entry, fmt.Errorf("decode %s : %v", entry.FullPath, err)
	}

	return entry, nil
}

func (store *AbstractSqlStore) DeleteEntry(ctx context.Context, fullpath util.FullPath) error {

	dir, name := fullpath.DirAndName()

	res, err := store.getTxOrDB(ctx).ExecContext(ctx, store.SqlDelete, util.HashStringToLong(dir), name, dir)
	if err != nil {
		return fmt.Errorf("delete %s: %s", fullpath, err)
	}

	_, err = res.RowsAffected()
	if err != nil {
		return fmt.Errorf("delete %s but no rows affected: %s", fullpath, err)
	}

	return nil
}

func (store *AbstractSqlStore) DeleteFolderChildren(ctx context.Context, fullpath util.FullPath) error {

	res, err := store.getTxOrDB(ctx).ExecContext(ctx, store.SqlDeleteFolderChildren, util.HashStringToLong(string(fullpath)), fullpath)
	if err != nil {
		return fmt.Errorf("deleteFolderChildren %s: %s", fullpath, err)
	}

	_, err = res.RowsAffected()
	if err != nil {
		return fmt.Errorf("deleteFolderChildren %s but no rows affected: %s", fullpath, err)
	}

	return nil
}

func (store *AbstractSqlStore) ListDirectoryPrefixedEntries(ctx context.Context, fullpath util.FullPath, startFileName string, inclusive bool, limit int, prefix string) (entries []*filer.Entry, err error) {
	sqlText := store.SqlListExclusive
	if inclusive {
		sqlText = store.SqlListInclusive
	}

	rows, err := store.getTxOrDB(ctx).QueryContext(ctx, sqlText, util.HashStringToLong(string(fullpath)), startFileName, string(fullpath), prefix, limit)
	if err != nil {
		return nil, fmt.Errorf("list %s : %v", fullpath, err)
	}
	defer rows.Close()

	for rows.Next() {
		var name string
		var data []byte
		if err = rows.Scan(&name, &data); err != nil {
			glog.V(0).Infof("scan %s : %v", fullpath, err)
			return nil, fmt.Errorf("scan %s: %v", fullpath, err)
		}

		entry := &filer.Entry{
			FullPath: util.NewFullPath(string(fullpath), name),
		}
		if err = entry.DecodeAttributesAndChunks(data); err != nil {
			glog.V(0).Infof("scan decode %s : %v", entry.FullPath, err)
			return nil, fmt.Errorf("scan decode %s : %v", entry.FullPath, err)
		}

		entries = append(entries, entry)
	}

	return entries, nil
}

func (store *AbstractSqlStore) ListDirectoryEntries(ctx context.Context, fullpath util.FullPath, startFileName string, inclusive bool, limit int) (entries []*filer.Entry, err error) {
	return store.ListDirectoryPrefixedEntries(ctx, fullpath, startFileName, inclusive, limit, "")
}

func (store *AbstractSqlStore) Shutdown() {
	store.DB.Close()
}
<|endoftext|>


<commit_before>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

// +build !linux

package diskmanager

import (
	"github.com/juju/juju/storage"
	"github.com/juju/juju/version"
)

var blockDeviceInUse = func(storage.BlockDevice) (bool, error) {
	panic("not supported")
}

func listBlockDevices() ([]storage.BlockDevice, error) {
	// Return an empty list each time.
	return nil, nil
}

func init() {
	logger.Infof(
		"block device support has not been implemented for %s",
		version.Current.OS,
	)
	DefaultListBlockDevices = listBlockDevices
}
<commit_msg>Fix compile failure for windows.<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

// +build !linux

package diskmanager

import (
	"runtime"

	"github.com/juju/juju/storage"
	"github.com/juju/juju/version"
)

var blockDeviceInUse = func(storage.BlockDevice) (bool, error) {
	panic("not supported")
}

func listBlockDevices() ([]storage.BlockDevice, error) {
	// Return an empty list each time.
	return nil, nil
}

func init() {
	logger.Infof(
		"block device support has not been implemented for %s",
		runtime.GOOS,
	)
	DefaultListBlockDevices = listBlockDevices
}
<|endoftext|>


<commit_before>// Copyright 2011 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.

// Package build gathers information about Go packages.
//
// Go Path
//
// The Go path is a list of directory trees containing Go source code.
// It is consulted to resolve imports that cannot be found in the standard
// Go tree.  The default path is the value of the GOPATH environment
// variable, interpreted as a path list appropriate to the operating system
// (on Unix, the variable is a colon-separated string;
// on Windows, a semicolon-separated string;
// on Plan 9, a list).
//
// Each directory listed in the Go path must have a prescribed structure:
//
// The src/ directory holds source code.  The path below 'src' determines
// the import path or executable name.
//
// The pkg/ directory holds installed package objects.
// As in the Go tree, each target operating system and
// architecture pair has its own subdirectory of pkg
// (pkg/GOOS_GOARCH).
//
// If DIR is a directory listed in the Go path, a package with
// source in DIR/src/foo/bar can be imported as "foo/bar" and
// has its compiled form installed to "DIR/pkg/GOOS_GOARCH/foo/bar.a"
// (or, for gccgo, "DIR/pkg/gccgo/foo/libbar.a").
//
// The bin/ directory holds compiled commands.
// Each command is named for its source directory, but only
// using the final element, not the entire path.  That is, the
// command with source in DIR/src/foo/quux is installed into
// DIR/bin/quux, not DIR/bin/foo/quux.  The foo/ is stripped
// so that you can add DIR/bin to your PATH to get at the
// installed commands.
//
// Here's an example directory layout:
//
//	GOPATH=/home/user/gocode
//
//	/home/user/gocode/
//	    src/
//	        foo/
//	            bar/               (go code in package bar)
//	                x.go
//	            quux/              (go code in package main)
//	                y.go
//	    bin/
//	        quux                   (installed command)
//	    pkg/
//	        linux_amd64/
//	            foo/
//	                bar.a          (installed package object)
//
// Build Constraints
//
// A build constraint is a line comment beginning with the directive +build
// that lists the conditions under which a file should be included in the package.
// Constraints may appear in any kind of source file (not just Go), but
// they must appear near the top of the file, preceded
// only by blank lines and other line comments.
//
// A build constraint is evaluated as the OR of space-separated options;
// each option evaluates as the AND of its comma-separated terms;
// and each term is an alphanumeric word or, preceded by !, its negation.
// That is, the build constraint:
//
//	// +build linux,386 darwin,!cgo
//
// corresponds to the boolean formula:
//
//	(linux AND 386) OR (darwin AND (NOT cgo))
//
// A file may have multiple build constraints. The overall constraint is the AND
// of the individual constraints. That is, the build constraints:
//
//	// +build linux darwin
//	// +build 386
//
// corresponds to the boolean formula:
//
//	(linux OR darwin) AND 386
//
// During a particular build, the following words are satisfied:
//
//	- the target operating system, as spelled by runtime.GOOS
//	- the target architecture, as spelled by runtime.GOARCH
//	- the compiler being used, currently either "gc" or "gccgo"
//	- "cgo", if ctxt.CgoEnabled is true
//	- any additional words listed in ctxt.BuildTags
//
// If a file's name, after stripping the extension and a possible _test suffix,
// matches *_GOOS, *_GOARCH, or *_GOOS_GOARCH for any known operating
// system and architecture values, then the file is considered to have an implicit
// build constraint requiring those terms.
//
// To keep a file from being considered for the build:
//
//	// +build ignore
//
// (any other unsatisfied word will work as well, but ``ignore'' is conventional.)
//
// To build a file only when using cgo, and only on Linux and OS X:
//
//	// +build linux,cgo darwin,cgo
//
// Such a file is usually paired with another file implementing the
// default functionality for other systems, which in this case would
// carry the constraint:
//
//	// +build !linux,!darwin !cgo
//
// Naming a file dns_windows.go will cause it to be included only when
// building the package for Windows; similarly, math_386.s will be included
// only when building the package for 32-bit x86.
//
package build
<commit_msg>go/build: document blank line required after build constraints<commit_after>// Copyright 2011 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.

// Package build gathers information about Go packages.
//
// Go Path
//
// The Go path is a list of directory trees containing Go source code.
// It is consulted to resolve imports that cannot be found in the standard
// Go tree.  The default path is the value of the GOPATH environment
// variable, interpreted as a path list appropriate to the operating system
// (on Unix, the variable is a colon-separated string;
// on Windows, a semicolon-separated string;
// on Plan 9, a list).
//
// Each directory listed in the Go path must have a prescribed structure:
//
// The src/ directory holds source code.  The path below 'src' determines
// the import path or executable name.
//
// The pkg/ directory holds installed package objects.
// As in the Go tree, each target operating system and
// architecture pair has its own subdirectory of pkg
// (pkg/GOOS_GOARCH).
//
// If DIR is a directory listed in the Go path, a package with
// source in DIR/src/foo/bar can be imported as "foo/bar" and
// has its compiled form installed to "DIR/pkg/GOOS_GOARCH/foo/bar.a"
// (or, for gccgo, "DIR/pkg/gccgo/foo/libbar.a").
//
// The bin/ directory holds compiled commands.
// Each command is named for its source directory, but only
// using the final element, not the entire path.  That is, the
// command with source in DIR/src/foo/quux is installed into
// DIR/bin/quux, not DIR/bin/foo/quux.  The foo/ is stripped
// so that you can add DIR/bin to your PATH to get at the
// installed commands.
//
// Here's an example directory layout:
//
//	GOPATH=/home/user/gocode
//
//	/home/user/gocode/
//	    src/
//	        foo/
//	            bar/               (go code in package bar)
//	                x.go
//	            quux/              (go code in package main)
//	                y.go
//	    bin/
//	        quux                   (installed command)
//	    pkg/
//	        linux_amd64/
//	            foo/
//	                bar.a          (installed package object)
//
// Build Constraints
//
// A build constraint is a line comment beginning with the directive +build
// that lists the conditions under which a file should be included in the package.
// Constraints may appear in any kind of source file (not just Go), but
// they must appear near the top of the file, preceded
// only by blank lines and other line comments.
//
// To distinguish build constraints from package documentation, a series of
// build constraints must be followed by a blank line.
//
// A build constraint is evaluated as the OR of space-separated options;
// each option evaluates as the AND of its comma-separated terms;
// and each term is an alphanumeric word or, preceded by !, its negation.
// That is, the build constraint:
//
//	// +build linux,386 darwin,!cgo
//
// corresponds to the boolean formula:
//
//	(linux AND 386) OR (darwin AND (NOT cgo))
//
// A file may have multiple build constraints. The overall constraint is the AND
// of the individual constraints. That is, the build constraints:
//
//	// +build linux darwin
//	// +build 386
//
// corresponds to the boolean formula:
//
//	(linux OR darwin) AND 386
//
// During a particular build, the following words are satisfied:
//
//	- the target operating system, as spelled by runtime.GOOS
//	- the target architecture, as spelled by runtime.GOARCH
//	- the compiler being used, currently either "gc" or "gccgo"
//	- "cgo", if ctxt.CgoEnabled is true
//	- any additional words listed in ctxt.BuildTags
//
// If a file's name, after stripping the extension and a possible _test suffix,
// matches *_GOOS, *_GOARCH, or *_GOOS_GOARCH for any known operating
// system and architecture values, then the file is considered to have an implicit
// build constraint requiring those terms.
//
// To keep a file from being considered for the build:
//
//	// +build ignore
//
// (any other unsatisfied word will work as well, but ``ignore'' is conventional.)
//
// To build a file only when using cgo, and only on Linux and OS X:
//
//	// +build linux,cgo darwin,cgo
//
// Such a file is usually paired with another file implementing the
// default functionality for other systems, which in this case would
// carry the constraint:
//
//	// +build !linux,!darwin !cgo
//
// Naming a file dns_windows.go will cause it to be included only when
// building the package for Windows; similarly, math_386.s will be included
// only when building the package for 32-bit x86.
//
package build
<|endoftext|>


<commit_before>// Copyright 2010 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.

package mime

import (
	"strings"
)

// isTSpecial returns true if rune is in 'tspecials' as defined by RFC
// 1531 and RFC 2045.
func isTSpecial(rune int) bool {
	return strings.IndexRune(`()<>@,;:\"/[]?=`, rune) != -1
}

// IsTokenChar returns true if rune is in 'token' as defined by RFC
// 1531 and RFC 2045.
func IsTokenChar(rune int) bool {
	// token := 1*<any (US-ASCII) CHAR except SPACE, CTLs,
	//             or tspecials>
	return rune > 0x20 && rune < 0x7f && !isTSpecial(rune)
}

// IsQText returns true if rune is in 'qtext' as defined by RFC 822.
func IsQText(rune int) bool {
	// CHAR        =  <any ASCII character>        ; (  0-177,  0.-127.)
	// qtext       =  <any CHAR excepting <">,     ; => may be folded
	//                "\" & CR, and including
	//                linear-white-space>
	switch rune {
	case int('"'), int('\\'), int('\r'):
		return false
	}
	return rune < 0x80
}
<commit_msg>mime: delete unnecessary constant conversions.<commit_after>// Copyright 2010 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.

package mime

import (
	"strings"
)

// isTSpecial returns true if rune is in 'tspecials' as defined by RFC
// 1531 and RFC 2045.
func isTSpecial(rune int) bool {
	return strings.IndexRune(`()<>@,;:\"/[]?=`, rune) != -1
}

// IsTokenChar returns true if rune is in 'token' as defined by RFC
// 1531 and RFC 2045.
func IsTokenChar(rune int) bool {
	// token := 1*<any (US-ASCII) CHAR except SPACE, CTLs,
	//             or tspecials>
	return rune > 0x20 && rune < 0x7f && !isTSpecial(rune)
}

// IsQText returns true if rune is in 'qtext' as defined by RFC 822.
func IsQText(rune int) bool {
	// CHAR        =  <any ASCII character>        ; (  0-177,  0.-127.)
	// qtext       =  <any CHAR excepting <">,     ; => may be folded
	//                "\" & CR, and including
	//                linear-white-space>
	switch rune {
	case '"', '\\', '\r':
		return false
	}
	return rune < 0x80
}
<|endoftext|>


<commit_before>package router

import (
	"bufio"
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"github.com/apex/log"
	"github.com/buger/jsonparser"
	"github.com/gin-gonic/gin"
	"github.com/mholt/archiver/v3"
	"github.com/pkg/errors"
	"github.com/pterodactyl/wings/api"
	"github.com/pterodactyl/wings/config"
	"github.com/pterodactyl/wings/installer"
	"github.com/pterodactyl/wings/router/tokens"
	"github.com/pterodactyl/wings/server"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

func getServerArchive(c *gin.Context) {
	auth := strings.SplitN(c.GetHeader("Authorization"), " ", 2)

	if len(auth) != 2 || auth[0] != "Bearer" {
		c.Header("WWW-Authenticate", "Bearer")
		c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
			"error": "The required authorization heads were not present in the request.",
		})
		return
	}

	token := tokens.TransferPayload{}
	if err := tokens.ParseToken([]byte(auth[1]), &token); err != nil {
		TrackedError(err).AbortWithServerError(c)
		return
	}

	if token.Subject != c.Param("server") {
		c.AbortWithStatusJSON(http.StatusForbidden, gin.H{
			"error": "( .. •˘___˘• .. )",
		})
		return
	}

	s := GetServer(c.Param("server"))

	st, err := s.Archiver.Stat()
	if err != nil {
		if !errors.Is(err, os.ErrNotExist) {
			TrackedServerError(err, s).SetMessage("failed to stat archive").AbortWithServerError(c)
			return
		}

		c.AbortWithStatus(http.StatusNotFound)
		return
	}

	checksum, err := s.Archiver.Checksum()
	if err != nil {
		TrackedServerError(err, s).SetMessage("failed to calculate checksum").AbortWithServerError(c)
		return
	}

	file, err := os.Open(s.Archiver.Path())
	if err != nil {
		tserr := TrackedServerError(err, s)
		if !os.IsNotExist(err) {
			tserr.SetMessage("failed to open archive for reading")
		} else {
			tserr.SetMessage("failed to open archive")
		}

		tserr.AbortWithServerError(c)
		return
	}
	defer file.Close()

	c.Header("X-Checksum", checksum)
	c.Header("X-Mime-Type", st.Mimetype)
	c.Header("Content-Length", strconv.Itoa(int(st.Info.Size())))
	c.Header("Content-Disposition", "attachment; filename="+s.Archiver.Name())
	c.Header("Content-Type", "application/octet-stream")

	bufio.NewReader(file).WriteTo(c.Writer)
}

func postServerArchive(c *gin.Context) {
	s := GetServer(c.Param("server"))

	go func(s *server.Server) {
		if err := s.Archiver.Archive(); err != nil {
			s.Log().WithField("error", err).Error("failed to get archive for server")
			return
		}

		s.Log().Debug("successfully created server archive, notifying panel")

		r := api.New()
		err := r.SendArchiveStatus(s.Id(), true)
		if err != nil {
			if !api.IsRequestError(err) {
				s.Log().WithField("error", err).Error("failed to notify panel of archive status")
				return
			}

			s.Log().WithField("error", err.Error()).Error("panel returned an error when sending the archive status")

			return
		}

		s.Log().Debug("successfully notified panel of archive status")
	}(s)

	c.Status(http.StatusAccepted)
}

func postTransfer(c *gin.Context) {
	buf := bytes.Buffer{}
	buf.ReadFrom(c.Request.Body)

	go func(data []byte) {
		serverID, _ := jsonparser.GetString(data, "server_id")
		url, _ := jsonparser.GetString(data, "url")
		token, _ := jsonparser.GetString(data, "token")

		l := log.WithField("server", serverID)
		// Create an http client with no timeout.
		client := &http.Client{Timeout: 0}

		hasError := true
		defer func() {
			if !hasError {
				return
			}

			l.Info("server transfer failed, notifying panel")
			err := api.New().SendTransferFailure(serverID)
			if err != nil {
				if !api.IsRequestError(err) {
					l.WithField("error", err).Error("failed to notify panel with transfer failure")
					return
				}

				l.WithField("error", err.Error()).Error("received error response from panel while notifying of transfer failure")
				return
			}

			l.Debug("notified panel of transfer failure")
		}()

		// Make a new GET request to the URL the panel gave us.
		req, err := http.NewRequest("GET", url, nil)
		if err != nil {
			log.WithField("error", err).Error("failed to create http request for archive transfer")
			return
		}

		// Add the authorization header.
		req.Header.Set("Authorization", token)

		// Execute the http request.
		res, err := client.Do(req)
		if err != nil {
			l.WithField("error", err).Error("failed to send archive http request")
			return
		}
		defer res.Body.Close()

		// Handle non-200 status codes.
		if res.StatusCode != 200 {
			_, err := ioutil.ReadAll(res.Body)
			if err != nil {
				l.WithField("error", err).WithField("status", res.StatusCode).Error("failed read transfer response body")
				return
			}

			l.WithField("error", err).WithField("status", res.StatusCode).Error("failed to request server archive")
			return
		}

		// Get the path to the archive.
		archivePath := filepath.Join(config.Get().System.ArchiveDirectory, serverID+".tar.gz")

		// Check if the archive already exists and delete it if it does.
		_, err = os.Stat(archivePath)
		if err != nil {
			if !os.IsNotExist(err) {
				l.WithField("error", err).Error("failed to stat archive file")
				return
			}
		} else {
			if err := os.Remove(archivePath); err != nil {
				l.WithField("error", err).Warn("failed to remove old archive file")
				return
			}
		}

		// Create the file.
		file, err := os.Create(archivePath)
		if err != nil {
			l.WithField("error", err).Error("failed to open archive on disk")
			return
		}

		// Copy the file.
		buf := make([]byte, 1024*4)
		_, err = io.CopyBuffer(file, res.Body, buf)
		if err != nil {
			l.WithField("error", err).Error("failed to copy archive file to disk")
			return
		}

		// Close the file so it can be opened to verify the checksum.
		if err := file.Close(); err != nil {
			l.WithField("error", err).Error("failed to close archive file")
			return
		}

		// Whenever the transfer fails or succeeds, delete the temporary transfer archive.
		defer func() {
			log.WithField("server", serverID).Debug("deleting temporary transfer archive..")
			if err := os.Remove(archivePath); err != nil && !os.IsNotExist(err) {
				l.WithFields(log.Fields{
					"server": serverID,
					"error":  err,
				}).Warn("failed to delete transfer archive")
			} else {
				l.WithField("server", serverID).Debug("deleted temporary transfer archive successfully")
			}
		}()

		l.WithField("server", serverID).Debug("server archive downloaded, computing checksum...")

		// Open the archive file for computing a checksum.
		file, err = os.Open(archivePath)
		if err != nil {
			l.WithField("error", err).Error("failed to open archive on disk")
			return
		}

		// Compute the sha256 checksum of the file.
		hash := sha256.New()
		buf = make([]byte, 1024*4)
		if _, err := io.CopyBuffer(hash, file, buf); err != nil {
			l.WithField("error", err).Error("failed to copy archive file for checksum verification")
			return
		}

		// Verify the two checksums.
		if hex.EncodeToString(hash.Sum(nil)) != res.Header.Get("X-Checksum") {
			l.Error("checksum verification failed for archive")
			return
		}

		// Close the file.
		if err := file.Close(); err != nil {
			l.WithField("error", err).Error("failed to close archive file after calculating checksum")
			return
		}

		l.Info("server archive transfer was successful")

		// Get the server data from the request.
		serverData, t, _, _ := jsonparser.Get(data, "server")
		if t != jsonparser.Object {
			l.Error("invalid server data passed in request")
			return
		}

		// Create a new server installer (note this does not execute the install script)
		i, err := installer.New(serverData)
		if err != nil {
			l.WithField("error", err).Error("failed to validate received server data")
			return
		}

		// Add the server to the collection.
		server.GetServers().Add(i.Server())

		// Create the server's environment (note this does not execute the install script)
		if err := i.Server().CreateEnvironment(); err != nil {
			l.WithField("error", err).Error("failed to create server environment")
			return
		}

		// Un-archive the archive, that sounds weird..
		if err := archiver.NewTarGz().Unarchive(archivePath, i.Server().Filesystem().Path()); err != nil {
			l.WithField("error", err).Error("failed to extract server archive")
			return
		}

		// We mark the process as being successful here as if we fail to send a transfer success,
		// then a transfer failure won't probably be successful either.
		//
		// It may be useful to retry sending the transfer success every so often just in case of a small
		// hiccup or the fix of whatever error causing the success request to fail.
		hasError = false

		// Notify the panel that the transfer succeeded.
		err = api.New().SendTransferSuccess(serverID)
		if err != nil {
			if !api.IsRequestError(err) {
				l.WithField("error", err).Error("failed to notify panel of transfer success")
				return
			}

			l.WithField("error", err.Error()).Error("panel responded with error after transfer success")
			return
		}

		l.WithField("server", serverID).Info("successfully notified panel of transfer success")
	}(buf.Bytes())

	c.Status(http.StatusAccepted)
}
<commit_msg>Notify panel of failed archive generation when transferring a server<commit_after>package router

import (
	"bufio"
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"github.com/apex/log"
	"github.com/buger/jsonparser"
	"github.com/gin-gonic/gin"
	"github.com/mholt/archiver/v3"
	"github.com/pkg/errors"
	"github.com/pterodactyl/wings/api"
	"github.com/pterodactyl/wings/config"
	"github.com/pterodactyl/wings/installer"
	"github.com/pterodactyl/wings/router/tokens"
	"github.com/pterodactyl/wings/server"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

func getServerArchive(c *gin.Context) {
	auth := strings.SplitN(c.GetHeader("Authorization"), " ", 2)

	if len(auth) != 2 || auth[0] != "Bearer" {
		c.Header("WWW-Authenticate", "Bearer")
		c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
			"error": "The required authorization heads were not present in the request.",
		})
		return
	}

	token := tokens.TransferPayload{}
	if err := tokens.ParseToken([]byte(auth[1]), &token); err != nil {
		TrackedError(err).AbortWithServerError(c)
		return
	}

	if token.Subject != c.Param("server") {
		c.AbortWithStatusJSON(http.StatusForbidden, gin.H{
			"error": "( .. •˘___˘• .. )",
		})
		return
	}

	s := GetServer(c.Param("server"))

	st, err := s.Archiver.Stat()
	if err != nil {
		if !errors.Is(err, os.ErrNotExist) {
			TrackedServerError(err, s).SetMessage("failed to stat archive").AbortWithServerError(c)
			return
		}

		c.AbortWithStatus(http.StatusNotFound)
		return
	}

	checksum, err := s.Archiver.Checksum()
	if err != nil {
		TrackedServerError(err, s).SetMessage("failed to calculate checksum").AbortWithServerError(c)
		return
	}

	file, err := os.Open(s.Archiver.Path())
	if err != nil {
		tserr := TrackedServerError(err, s)
		if !os.IsNotExist(err) {
			tserr.SetMessage("failed to open archive for reading")
		} else {
			tserr.SetMessage("failed to open archive")
		}

		tserr.AbortWithServerError(c)
		return
	}
	defer file.Close()

	c.Header("X-Checksum", checksum)
	c.Header("X-Mime-Type", st.Mimetype)
	c.Header("Content-Length", strconv.Itoa(int(st.Info.Size())))
	c.Header("Content-Disposition", "attachment; filename="+s.Archiver.Name())
	c.Header("Content-Type", "application/octet-stream")

	bufio.NewReader(file).WriteTo(c.Writer)
}

func postServerArchive(c *gin.Context) {
	s := GetServer(c.Param("server"))

	go func(s *server.Server) {
		r := api.New()

		// Attempt to get an archive of the server.  This **WILL NOT** modify the source files of a server,
		// this process is 100% safe and will not corrupt a server's files if it fails.
		if err := s.Archiver.Archive(); err != nil {
			s.Log().WithField("error", err).Error("failed to get archive for server")

			if err := r.SendArchiveStatus(s.Id(), false); err != nil {
				if !api.IsRequestError(err) {
					s.Log().WithField("error", err).Error("failed to notify panel of failed archive status")
					return
				}

				s.Log().WithField("error", err.Error()).Error("panel returned an error when notifying it of a failed archive status")
				return
			}

			s.Log().Info("successfully notified panel of failed archive status")
			return
		}

		s.Log().Debug("successfully created server archive, notifying panel")

		if err := r.SendArchiveStatus(s.Id(), true); err != nil {
			if !api.IsRequestError(err) {
				s.Log().WithField("error", err).Error("failed to notify panel of successful archive status")
				return
			}

			s.Log().WithField("error", err.Error()).Error("panel returned an error when notifying it of a successful archive status")
			return
		}

		s.Log().Info("successfully notified panel of successful archive status")
	}(s)

	c.Status(http.StatusAccepted)
}

func postTransfer(c *gin.Context) {
	buf := bytes.Buffer{}
	buf.ReadFrom(c.Request.Body)

	go func(data []byte) {
		serverID, _ := jsonparser.GetString(data, "server_id")
		url, _ := jsonparser.GetString(data, "url")
		token, _ := jsonparser.GetString(data, "token")

		l := log.WithField("server", serverID)
		// Create an http client with no timeout.
		client := &http.Client{Timeout: 0}

		hasError := true
		defer func() {
			if !hasError {
				return
			}

			l.Info("server transfer failed, notifying panel")
			err := api.New().SendTransferFailure(serverID)
			if err != nil {
				if !api.IsRequestError(err) {
					l.WithField("error", err).Error("failed to notify panel with transfer failure")
					return
				}

				l.WithField("error", err.Error()).Error("received error response from panel while notifying of transfer failure")
				return
			}

			l.Debug("notified panel of transfer failure")
		}()

		// Make a new GET request to the URL the panel gave us.
		req, err := http.NewRequest("GET", url, nil)
		if err != nil {
			log.WithField("error", err).Error("failed to create http request for archive transfer")
			return
		}

		// Add the authorization header.
		req.Header.Set("Authorization", token)

		// Execute the http request.
		res, err := client.Do(req)
		if err != nil {
			l.WithField("error", err).Error("failed to send archive http request")
			return
		}
		defer res.Body.Close()

		// Handle non-200 status codes.
		if res.StatusCode != 200 {
			_, err := ioutil.ReadAll(res.Body)
			if err != nil {
				l.WithField("error", err).WithField("status", res.StatusCode).Error("failed read transfer response body")
				return
			}

			l.WithField("error", err).WithField("status", res.StatusCode).Error("failed to request server archive")
			return
		}

		// Get the path to the archive.
		archivePath := filepath.Join(config.Get().System.ArchiveDirectory, serverID+".tar.gz")

		// Check if the archive already exists and delete it if it does.
		_, err = os.Stat(archivePath)
		if err != nil {
			if !os.IsNotExist(err) {
				l.WithField("error", err).Error("failed to stat archive file")
				return
			}
		} else {
			if err := os.Remove(archivePath); err != nil {
				l.WithField("error", err).Warn("failed to remove old archive file")
				return
			}
		}

		// Create the file.
		file, err := os.Create(archivePath)
		if err != nil {
			l.WithField("error", err).Error("failed to open archive on disk")
			return
		}

		// Copy the file.
		buf := make([]byte, 1024*4)
		_, err = io.CopyBuffer(file, res.Body, buf)
		if err != nil {
			l.WithField("error", err).Error("failed to copy archive file to disk")
			return
		}

		// Close the file so it can be opened to verify the checksum.
		if err := file.Close(); err != nil {
			l.WithField("error", err).Error("failed to close archive file")
			return
		}

		// Whenever the transfer fails or succeeds, delete the temporary transfer archive.
		defer func() {
			log.WithField("server", serverID).Debug("deleting temporary transfer archive..")
			if err := os.Remove(archivePath); err != nil && !os.IsNotExist(err) {
				l.WithFields(log.Fields{
					"server": serverID,
					"error":  err,
				}).Warn("failed to delete transfer archive")
			} else {
				l.WithField("server", serverID).Debug("deleted temporary transfer archive successfully")
			}
		}()

		l.WithField("server", serverID).Debug("server archive downloaded, computing checksum...")

		// Open the archive file for computing a checksum.
		file, err = os.Open(archivePath)
		if err != nil {
			l.WithField("error", err).Error("failed to open archive on disk")
			return
		}

		// Compute the sha256 checksum of the file.
		hash := sha256.New()
		buf = make([]byte, 1024*4)
		if _, err := io.CopyBuffer(hash, file, buf); err != nil {
			l.WithField("error", err).Error("failed to copy archive file for checksum verification")
			return
		}

		// Verify the two checksums.
		if hex.EncodeToString(hash.Sum(nil)) != res.Header.Get("X-Checksum") {
			l.Error("checksum verification failed for archive")
			return
		}

		// Close the file.
		if err := file.Close(); err != nil {
			l.WithField("error", err).Error("failed to close archive file after calculating checksum")
			return
		}

		l.Info("server archive transfer was successful")

		// Get the server data from the request.
		serverData, t, _, _ := jsonparser.Get(data, "server")
		if t != jsonparser.Object {
			l.Error("invalid server data passed in request")
			return
		}

		// Create a new server installer (note this does not execute the install script)
		i, err := installer.New(serverData)
		if err != nil {
			l.WithField("error", err).Error("failed to validate received server data")
			return
		}

		// Add the server to the collection.
		server.GetServers().Add(i.Server())

		// Create the server's environment (note this does not execute the install script)
		if err := i.Server().CreateEnvironment(); err != nil {
			l.WithField("error", err).Error("failed to create server environment")
			return
		}

		// Un-archive the archive, that sounds weird..
		if err := archiver.NewTarGz().Unarchive(archivePath, i.Server().Filesystem().Path()); err != nil {
			l.WithField("error", err).Error("failed to extract server archive")
			return
		}

		// We mark the process as being successful here as if we fail to send a transfer success,
		// then a transfer failure won't probably be successful either.
		//
		// It may be useful to retry sending the transfer success every so often just in case of a small
		// hiccup or the fix of whatever error causing the success request to fail.
		hasError = false

		// Notify the panel that the transfer succeeded.
		err = api.New().SendTransferSuccess(serverID)
		if err != nil {
			if !api.IsRequestError(err) {
				l.WithField("error", err).Error("failed to notify panel of transfer success")
				return
			}

			l.WithField("error", err.Error()).Error("panel responded with error after transfer success")
			return
		}

		l.WithField("server", serverID).Info("successfully notified panel of transfer success")
	}(buf.Bytes())

	c.Status(http.StatusAccepted)
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

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 onlineddl

import (
	"flag"
	"fmt"
	"os"
	"path"
	"regexp"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"vitess.io/vitess/go/test/endtoend/cluster"
)

var (
	clusterInstance       *cluster.LocalProcessCluster
	hostname              = "localhost"
	keyspaceName          = "ks"
	cell                  = "zone1"
	schemaChangeDirectory = ""
	totalTableCount       = 4
	createTable           = `
		CREATE TABLE %s (
		id BIGINT(20) not NULL,
		msg varchar(64),
		PRIMARY KEY (id)
		) ENGINE=InnoDB;`
	alterTable = `
		ALTER WITH 'gh-ost' TABLE %s
		MODIFY id BIGINT UNSIGNED NOT NULL,
		ADD COLUMN ghost_col INT NOT NULL,
		ADD INDEX idx_msg(msg)`
	statusCompleteRegexp = regexp.MustCompile(`\bcomplete\b`)
)

func TestMain(m *testing.M) {
	defer cluster.PanicHandler(nil)
	flag.Parse()

	exitcode, err := func() (int, error) {
		clusterInstance = cluster.NewCluster(cell, hostname)
		schemaChangeDirectory = path.Join("/tmp", fmt.Sprintf("schema_change_dir_%d", clusterInstance.GetAndReserveTabletUID()))
		defer os.RemoveAll(schemaChangeDirectory)
		defer clusterInstance.Teardown()

		if _, err := os.Stat(schemaChangeDirectory); os.IsNotExist(err) {
			_ = os.Mkdir(schemaChangeDirectory, 0700)
		}

		clusterInstance.VtctldExtraArgs = []string{
			"-schema_change_dir", schemaChangeDirectory,
			"-schema_change_controller", "local",
			"-schema_change_check_interval", "1"}

		if err := clusterInstance.StartTopo(); err != nil {
			return 1, err
		}

		// Start keyspace
		keyspace := &cluster.Keyspace{
			Name: keyspaceName,
		}

		if err := clusterInstance.StartUnshardedKeyspace(*keyspace, 2, true); err != nil {
			return 1, err
		}
		if err := clusterInstance.StartKeyspace(*keyspace, []string{"1"}, 1, false); err != nil {
			return 1, err
		}
		return m.Run(), nil
	}()
	if err != nil {
		fmt.Printf("%v\n", err)
		os.Exit(1)
	} else {
		os.Exit(exitcode)
	}

}

func TestSchemaChange(t *testing.T) {
	defer cluster.PanicHandler(t)
	testWithInitialSchema(t)
	testWithAlterSchema(t)
}

func testWithInitialSchema(t *testing.T) {
	// Create 4 tables
	var sqlQuery = "" //nolint
	for i := 0; i < totalTableCount; i++ {
		sqlQuery = fmt.Sprintf(createTable, fmt.Sprintf("vt_onlineddl_test_%02d", i))
		err := clusterInstance.VtctlclientProcess.ApplySchema(keyspaceName, sqlQuery)
		require.Nil(t, err)
	}

	// Check if 4 tables are created
	checkTables(t, totalTableCount)
	checkTables(t, totalTableCount)
}

// testWithAlterSchema if we alter schema and then apply, the resultant schema should match across shards
func testWithAlterSchema(t *testing.T) {
	tableName := fmt.Sprintf("vt_onlineddl_test_%02d", 3)
	sqlQuery := fmt.Sprintf(alterTable, tableName)
	uuid, err := clusterInstance.VtctlclientProcess.ApplySchemaWithOutput(keyspaceName, sqlQuery)
	require.Nil(t, err)
	uuid = strings.TrimSpace(uuid)
	require.NotEmpty(t, uuid)
	// Migration is asynchronous. Give it some time.
	time.Sleep(time.Second * 30)
	checkRecentMigrations(t, tableName, uuid)
	checkMigratedTable(t, tableName)
}

// checkTables checks the number of tables in the first two shards.
func checkTables(t *testing.T, count int) {
	checkTablesCount(t, clusterInstance.Keyspaces[0].Shards[0].Vttablets[0], count)
	checkTablesCount(t, clusterInstance.Keyspaces[0].Shards[1].Vttablets[0], count)
}

// checkTablesCount checks the number of tables in the given tablet
func checkTablesCount(t *testing.T, tablet *cluster.Vttablet, count int) {
	queryResult, err := tablet.VttabletProcess.QueryTablet("show tables;", keyspaceName, true)
	require.Nil(t, err)
	assert.Equal(t, len(queryResult.Rows), count)
}

func checkRecentMigrations(t *testing.T, tableName, uuid string) {
	result, err := clusterInstance.VtctlclientProcess.OnlineDDLShowRecent(keyspaceName)
	assert.NoError(t, err)
	assert.Equal(t, len(clusterInstance.Keyspaces[0].Shards), strings.Count(result, tableName))
	assert.Equal(t, len(clusterInstance.Keyspaces[0].Shards), strings.Count(result, uuid))
	// The word "complete" appears in the column `completed_timestamp`. So we use a regexp to
	// ensure we match exact full word
	m := statusCompleteRegexp.FindAllString(result, -1)
	assert.Equal(t, len(clusterInstance.Keyspaces[0].Shards), len(m))
}

// checkMigratedTables checks the CREATE STATEMENT of a table after migration
func checkMigratedTable(t *testing.T, tableName string) {
	expect := "ghost_col"
	checkTableCreateContains(t, clusterInstance.Keyspaces[0].Shards[0].Vttablets[0], tableName, expect)
}

// checkTableCreateContains checks if table's CREATE TABLE statement contains a given test
func checkTableCreateContains(t *testing.T, tablet *cluster.Vttablet, tableName string, expect string) {
	queryResult, err := tablet.VttabletProcess.QueryTablet(fmt.Sprintf("show create table %s;", tableName), keyspaceName, true)
	require.Nil(t, err)

	assert.Equal(t, len(queryResult.Rows), 1)
	assert.Equal(t, len(queryResult.Rows[0]), 2) // table name, create statement
	createStatement := queryResult.Rows[0][1].ToString()
	assert.True(t, strings.Contains(createStatement, expect))
}
<commit_msg>expect a failed migration<commit_after>/*
Copyright 2019 The Vitess Authors.

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 onlineddl

import (
	"flag"
	"fmt"
	"os"
	"path"
	"regexp"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"vitess.io/vitess/go/test/endtoend/cluster"
)

var (
	clusterInstance       *cluster.LocalProcessCluster
	hostname              = "localhost"
	keyspaceName          = "ks"
	cell                  = "zone1"
	schemaChangeDirectory = ""
	totalTableCount       = 4
	createTable           = `
		CREATE TABLE %s (
		id BIGINT(20) not NULL,
		msg varchar(64),
		PRIMARY KEY (id)
		) ENGINE=InnoDB;`
	alterTableStatament = `
		ALTER WITH 'gh-ost' TABLE %s
		MODIFY id BIGINT UNSIGNED NOT NULL,
		ADD COLUMN ghost_col INT NOT NULL,
		ADD INDEX idx_msg(msg)`
	statusCompleteRegexp = regexp.MustCompile(`\bcomplete\b`)
	statusFailedRegexp   = regexp.MustCompile(`\bfailed\b`)
)

func TestMain(m *testing.M) {
	defer cluster.PanicHandler(nil)
	flag.Parse()

	exitcode, err := func() (int, error) {
		clusterInstance = cluster.NewCluster(cell, hostname)
		schemaChangeDirectory = path.Join("/tmp", fmt.Sprintf("schema_change_dir_%d", clusterInstance.GetAndReserveTabletUID()))
		defer os.RemoveAll(schemaChangeDirectory)
		defer clusterInstance.Teardown()

		if _, err := os.Stat(schemaChangeDirectory); os.IsNotExist(err) {
			_ = os.Mkdir(schemaChangeDirectory, 0700)
		}

		clusterInstance.VtctldExtraArgs = []string{
			"-schema_change_dir", schemaChangeDirectory,
			"-schema_change_controller", "local",
			"-schema_change_check_interval", "1"}

		if err := clusterInstance.StartTopo(); err != nil {
			return 1, err
		}

		// Start keyspace
		keyspace := &cluster.Keyspace{
			Name: keyspaceName,
		}

		if err := clusterInstance.StartUnshardedKeyspace(*keyspace, 2, true); err != nil {
			return 1, err
		}
		if err := clusterInstance.StartKeyspace(*keyspace, []string{"1"}, 1, false); err != nil {
			return 1, err
		}
		return m.Run(), nil
	}()
	if err != nil {
		fmt.Printf("%v\n", err)
		os.Exit(1)
	} else {
		os.Exit(exitcode)
	}

}

func TestSchemaChange(t *testing.T) {
	defer cluster.PanicHandler(t)
	assert.Equal(t, 2, len(clusterInstance.Keyspaces[0].Shards))
	testWithInitialSchema(t)
	testWithValidAlterSchema(t, true)
	testWithValidAlterSchema(t, false)
}

func testWithInitialSchema(t *testing.T) {
	// Create 4 tables
	var sqlQuery = "" //nolint
	for i := 0; i < totalTableCount; i++ {
		sqlQuery = fmt.Sprintf(createTable, fmt.Sprintf("vt_onlineddl_test_%02d", i))
		err := clusterInstance.VtctlclientProcess.ApplySchema(keyspaceName, sqlQuery)
		require.Nil(t, err)
	}

	// Check if 4 tables are created
	checkTables(t, totalTableCount)
}

// testWithAlterSchema if we alter schema and then apply, the resultant schema should match across shards
func testWithValidAlterSchema(t *testing.T, expectSuccess bool) {
	tableName := fmt.Sprintf("vt_onlineddl_test_%02d", 3)
	sqlQuery := fmt.Sprintf(alterTableStatament, tableName)
	uuid, err := clusterInstance.VtctlclientProcess.ApplySchemaWithOutput(keyspaceName, sqlQuery)
	require.Nil(t, err)
	uuid = strings.TrimSpace(uuid)
	require.NotEmpty(t, uuid)
	// Migration is asynchronous. Give it some time.
	time.Sleep(time.Second * 30)
	if expectSuccess {
		checkRecentMigrations(t, uuid, statusCompleteRegexp)
	} else {
		checkRecentMigrations(t, uuid, statusFailedRegexp)
	}
	checkMigratedTable(t, tableName)
}

// checkTables checks the number of tables in the first two shards.
func checkTables(t *testing.T, count int) {
	for i := range clusterInstance.Keyspaces[0].Shards {
		checkTablesCount(t, clusterInstance.Keyspaces[0].Shards[i].Vttablets[0], count)
	}
}

// checkTablesCount checks the number of tables in the given tablet
func checkTablesCount(t *testing.T, tablet *cluster.Vttablet, count int) {
	queryResult, err := tablet.VttabletProcess.QueryTablet("show tables;", keyspaceName, true)
	require.Nil(t, err)
	assert.Equal(t, len(queryResult.Rows), count)
}

func checkRecentMigrations(t *testing.T, uuid string, expectStatusRegexp *regexp.Regexp) {
	result, err := clusterInstance.VtctlclientProcess.OnlineDDLShowRecent(keyspaceName)
	assert.NoError(t, err)
	assert.Equal(t, len(clusterInstance.Keyspaces[0].Shards), strings.Count(result, uuid))
	// The word "complete" appears in the column `completed_timestamp`. So we use a regexp to
	// ensure we match exact full word
	m := expectStatusRegexp.FindAllString(result, -1)
	assert.Equal(t, len(clusterInstance.Keyspaces[0].Shards), len(m))
}

// checkMigratedTables checks the CREATE STATEMENT of a table after migration
func checkMigratedTable(t *testing.T, tableName string) {
	expect := "ghost_col"
	for i := range clusterInstance.Keyspaces[0].Shards {
		createStatement := getCreateTableStatement(t, clusterInstance.Keyspaces[0].Shards[i].Vttablets[0], tableName)
		assert.True(t, strings.Contains(createStatement, expect))
	}
}

// getCreateTableStatement returns the CREATE TABLE statement for a given table
func getCreateTableStatement(t *testing.T, tablet *cluster.Vttablet, tableName string) (statement string) {
	queryResult, err := tablet.VttabletProcess.QueryTablet(fmt.Sprintf("show create table %s;", tableName), keyspaceName, true)
	require.Nil(t, err)

	assert.Equal(t, len(queryResult.Rows), 1)
	assert.Equal(t, len(queryResult.Rows[0]), 2) // table name, create statement
	statement = queryResult.Rows[0][1].ToString()
	return statement
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

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 onlineddl

import (
	"flag"
	"fmt"
	"os"
	"path"
	"regexp"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"vitess.io/vitess/go/test/endtoend/cluster"
)

var (
	clusterInstance       *cluster.LocalProcessCluster
	hostname              = "localhost"
	keyspaceName          = "ks"
	cell                  = "zone1"
	schemaChangeDirectory = ""
	totalTableCount       = 4
	createTable           = `
		CREATE TABLE %s (
		id BIGINT(20) not NULL,
		msg varchar(64),
		PRIMARY KEY (id)
		) ENGINE=InnoDB;`
	// The following statement is valid
	alterTableSuccessfulStatament = `
		ALTER WITH 'gh-ost' TABLE %s
		MODIFY id BIGINT UNSIGNED NOT NULL,
		ADD COLUMN ghost_col INT NOT NULL,
		ADD INDEX idx_msg(msg)`
	// The following statement will fail because gh-ost requires some shared unique key
	alterTableFailedStatament = `
		ALTER WITH 'gh-ost' TABLE %s
		DROP PRIMARY KEY`
	statusCompleteRegexp = regexp.MustCompile(`\bcomplete\b`)
	statusFailedRegexp   = regexp.MustCompile(`\bfailed\b`)
)

func TestMain(m *testing.M) {
	defer cluster.PanicHandler(nil)
	flag.Parse()

	exitcode, err := func() (int, error) {
		clusterInstance = cluster.NewCluster(cell, hostname)
		schemaChangeDirectory = path.Join("/tmp", fmt.Sprintf("schema_change_dir_%d", clusterInstance.GetAndReserveTabletUID()))
		defer os.RemoveAll(schemaChangeDirectory)
		defer clusterInstance.Teardown()

		if _, err := os.Stat(schemaChangeDirectory); os.IsNotExist(err) {
			_ = os.Mkdir(schemaChangeDirectory, 0700)
		}

		clusterInstance.VtctldExtraArgs = []string{
			"-schema_change_dir", schemaChangeDirectory,
			"-schema_change_controller", "local",
			"-schema_change_check_interval", "1"}

		if err := clusterInstance.StartTopo(); err != nil {
			return 1, err
		}

		// Start keyspace
		keyspace := &cluster.Keyspace{
			Name: keyspaceName,
		}

		if err := clusterInstance.StartUnshardedKeyspace(*keyspace, 2, true); err != nil {
			return 1, err
		}
		if err := clusterInstance.StartKeyspace(*keyspace, []string{"1"}, 1, false); err != nil {
			return 1, err
		}
		return m.Run(), nil
	}()
	if err != nil {
		fmt.Printf("%v\n", err)
		os.Exit(1)
	} else {
		os.Exit(exitcode)
	}

}

func TestSchemaChange(t *testing.T) {
	defer cluster.PanicHandler(t)
	assert.Equal(t, 2, len(clusterInstance.Keyspaces[0].Shards))
	testWithInitialSchema(t)
	// Expect first migration to complete
	testWithAlterSchema(t, alterTableSuccessfulStatament, true)
	// Expect 2nd invocation of same migration to fail
	testWithAlterSchema(t, alterTableFailedStatament, false)
}

func testWithInitialSchema(t *testing.T) {
	// Create 4 tables
	var sqlQuery = "" //nolint
	for i := 0; i < totalTableCount; i++ {
		sqlQuery = fmt.Sprintf(createTable, fmt.Sprintf("vt_onlineddl_test_%02d", i))
		err := clusterInstance.VtctlclientProcess.ApplySchema(keyspaceName, sqlQuery)
		require.Nil(t, err)
	}

	// Check if 4 tables are created
	checkTables(t, totalTableCount)
}

// testWithAlterSchema if we alter schema and then apply, the resultant schema should match across shards
func testWithAlterSchema(t *testing.T, alterStatement string, expectSuccess bool) {
	tableName := fmt.Sprintf("vt_onlineddl_test_%02d", 3)
	sqlQuery := fmt.Sprintf(alterStatement, tableName)
	uuid, err := clusterInstance.VtctlclientProcess.ApplySchemaWithOutput(keyspaceName, sqlQuery)
	require.Nil(t, err)
	uuid = strings.TrimSpace(uuid)
	require.NotEmpty(t, uuid)
	// Migration is asynchronous. Give it some time.
	time.Sleep(time.Second * 30)
	if expectSuccess {
		checkRecentMigrations(t, uuid, statusCompleteRegexp)
	} else {
		checkRecentMigrations(t, uuid, statusFailedRegexp)
	}
	checkMigratedTable(t, tableName)
	checkCancelMigration(t, uuid)
	// retry request should fail for successful migrations, and should succeed for failed migrations
	checkRetryMigration(t, uuid, !expectSuccess)
}

// checkTables checks the number of tables in the first two shards.
func checkTables(t *testing.T, count int) {
	for i := range clusterInstance.Keyspaces[0].Shards {
		checkTablesCount(t, clusterInstance.Keyspaces[0].Shards[i].Vttablets[0], count)
	}
}

// checkTablesCount checks the number of tables in the given tablet
func checkTablesCount(t *testing.T, tablet *cluster.Vttablet, count int) {
	queryResult, err := tablet.VttabletProcess.QueryTablet("show tables;", keyspaceName, true)
	require.Nil(t, err)
	assert.Equal(t, len(queryResult.Rows), count)
}

func checkRecentMigrations(t *testing.T, uuid string, expectStatusRegexp *regexp.Regexp) {
	result, err := clusterInstance.VtctlclientProcess.OnlineDDLShowRecent(keyspaceName)
	assert.NoError(t, err)
	fmt.Println("# 'vtctlclient OnlineDDL show recent' output (for debug purposes):")
	fmt.Println(result)
	assert.Equal(t, len(clusterInstance.Keyspaces[0].Shards), strings.Count(result, uuid))
	// The word "complete" appears in the column `completed_timestamp`. So we use a regexp to
	// ensure we match exact full word
	m := expectStatusRegexp.FindAllString(result, -1)
	assert.Equal(t, len(clusterInstance.Keyspaces[0].Shards), len(m))
}

// checkCancelMigration attempts to cancel a migration, and expects rejection
func checkCancelMigration(t *testing.T, uuid string) {
	result, err := clusterInstance.VtctlclientProcess.OnlineDDLCancelMigration(keyspaceName, uuid)
	assert.NoError(t, err)
	// The migration has either been complete or failed. We can't cancel it. Expect "zero" response from all tablets
	m := regexp.MustCompile("\b0\b").FindAllString(result, -1)
	assert.Equal(t, len(clusterInstance.Keyspaces[0].Shards), len(m))
}

// checkRetryMigration attempts to retry a migration, and expects rejection
func checkRetryMigration(t *testing.T, uuid string, expectSuccess bool) {
	result, err := clusterInstance.VtctlclientProcess.OnlineDDLRetryMigration(keyspaceName, uuid)
	assert.NoError(t, err)
	// The migration has either been complete or failed. We can't cancel it. Expect "zero" response from all tablets
	var r *regexp.Regexp
	if expectSuccess {
		r = regexp.MustCompile("\b1\b")
	} else {
		r = regexp.MustCompile("\b0\b")
	}
	m := r.FindAllString(result, -1)
	assert.Equal(t, len(clusterInstance.Keyspaces[0].Shards), len(m))
}

// checkMigratedTables checks the CREATE STATEMENT of a table after migration
func checkMigratedTable(t *testing.T, tableName string) {
	expect := "ghost_col"
	for i := range clusterInstance.Keyspaces[0].Shards {
		createStatement := getCreateTableStatement(t, clusterInstance.Keyspaces[0].Shards[i].Vttablets[0], tableName)
		assert.True(t, strings.Contains(createStatement, expect))
	}
}

// getCreateTableStatement returns the CREATE TABLE statement for a given table
func getCreateTableStatement(t *testing.T, tablet *cluster.Vttablet, tableName string) (statement string) {
	queryResult, err := tablet.VttabletProcess.QueryTablet(fmt.Sprintf("show create table %s;", tableName), keyspaceName, true)
	require.Nil(t, err)

	assert.Equal(t, len(queryResult.Rows), 1)
	assert.Equal(t, len(queryResult.Rows[0]), 2) // table name, create statement
	statement = queryResult.Rows[0][1].ToString()
	return statement
}
<commit_msg>debug info<commit_after>/*
Copyright 2019 The Vitess Authors.

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 onlineddl

import (
	"flag"
	"fmt"
	"os"
	"path"
	"regexp"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"vitess.io/vitess/go/test/endtoend/cluster"
)

var (
	clusterInstance       *cluster.LocalProcessCluster
	hostname              = "localhost"
	keyspaceName          = "ks"
	cell                  = "zone1"
	schemaChangeDirectory = ""
	totalTableCount       = 4
	createTable           = `
		CREATE TABLE %s (
		id BIGINT(20) not NULL,
		msg varchar(64),
		PRIMARY KEY (id)
		) ENGINE=InnoDB;`
	// The following statement is valid
	alterTableSuccessfulStatament = `
		ALTER WITH 'gh-ost' TABLE %s
		MODIFY id BIGINT UNSIGNED NOT NULL,
		ADD COLUMN ghost_col INT NOT NULL,
		ADD INDEX idx_msg(msg)`
	// The following statement will fail because gh-ost requires some shared unique key
	alterTableFailedStatament = `
		ALTER WITH 'gh-ost' TABLE %s
		DROP PRIMARY KEY`
	statusCompleteRegexp = regexp.MustCompile(`\bcomplete\b`)
	statusFailedRegexp   = regexp.MustCompile(`\bfailed\b`)
)

func TestMain(m *testing.M) {
	defer cluster.PanicHandler(nil)
	flag.Parse()

	exitcode, err := func() (int, error) {
		clusterInstance = cluster.NewCluster(cell, hostname)
		schemaChangeDirectory = path.Join("/tmp", fmt.Sprintf("schema_change_dir_%d", clusterInstance.GetAndReserveTabletUID()))
		defer os.RemoveAll(schemaChangeDirectory)
		defer clusterInstance.Teardown()

		if _, err := os.Stat(schemaChangeDirectory); os.IsNotExist(err) {
			_ = os.Mkdir(schemaChangeDirectory, 0700)
		}

		clusterInstance.VtctldExtraArgs = []string{
			"-schema_change_dir", schemaChangeDirectory,
			"-schema_change_controller", "local",
			"-schema_change_check_interval", "1"}

		if err := clusterInstance.StartTopo(); err != nil {
			return 1, err
		}

		// Start keyspace
		keyspace := &cluster.Keyspace{
			Name: keyspaceName,
		}

		if err := clusterInstance.StartUnshardedKeyspace(*keyspace, 2, true); err != nil {
			return 1, err
		}
		if err := clusterInstance.StartKeyspace(*keyspace, []string{"1"}, 1, false); err != nil {
			return 1, err
		}
		return m.Run(), nil
	}()
	if err != nil {
		fmt.Printf("%v\n", err)
		os.Exit(1)
	} else {
		os.Exit(exitcode)
	}

}

func TestSchemaChange(t *testing.T) {
	defer cluster.PanicHandler(t)
	assert.Equal(t, 2, len(clusterInstance.Keyspaces[0].Shards))
	testWithInitialSchema(t)
	// Expect first migration to complete
	testWithAlterSchema(t, alterTableSuccessfulStatament, true)
	// Expect 2nd invocation of same migration to fail
	testWithAlterSchema(t, alterTableFailedStatament, false)
}

func testWithInitialSchema(t *testing.T) {
	// Create 4 tables
	var sqlQuery = "" //nolint
	for i := 0; i < totalTableCount; i++ {
		sqlQuery = fmt.Sprintf(createTable, fmt.Sprintf("vt_onlineddl_test_%02d", i))
		err := clusterInstance.VtctlclientProcess.ApplySchema(keyspaceName, sqlQuery)
		require.Nil(t, err)
	}

	// Check if 4 tables are created
	checkTables(t, totalTableCount)
}

// testWithAlterSchema if we alter schema and then apply, the resultant schema should match across shards
func testWithAlterSchema(t *testing.T, alterStatement string, expectSuccess bool) {
	tableName := fmt.Sprintf("vt_onlineddl_test_%02d", 3)
	sqlQuery := fmt.Sprintf(alterStatement, tableName)
	uuid, err := clusterInstance.VtctlclientProcess.ApplySchemaWithOutput(keyspaceName, sqlQuery)
	require.Nil(t, err)
	uuid = strings.TrimSpace(uuid)
	require.NotEmpty(t, uuid)
	// Migration is asynchronous. Give it some time.
	time.Sleep(time.Second * 30)
	if expectSuccess {
		checkRecentMigrations(t, uuid, statusCompleteRegexp)
	} else {
		checkRecentMigrations(t, uuid, statusFailedRegexp)
	}
	checkMigratedTable(t, tableName)
	checkCancelMigration(t, uuid)
	// retry request should fail for successful migrations, and should succeed for failed migrations
	checkRetryMigration(t, uuid, !expectSuccess)
}

// checkTables checks the number of tables in the first two shards.
func checkTables(t *testing.T, count int) {
	for i := range clusterInstance.Keyspaces[0].Shards {
		checkTablesCount(t, clusterInstance.Keyspaces[0].Shards[i].Vttablets[0], count)
	}
}

// checkTablesCount checks the number of tables in the given tablet
func checkTablesCount(t *testing.T, tablet *cluster.Vttablet, count int) {
	queryResult, err := tablet.VttabletProcess.QueryTablet("show tables;", keyspaceName, true)
	require.Nil(t, err)
	assert.Equal(t, len(queryResult.Rows), count)
}

// checkRecentMigrations checks 'OnlineDDL <keyspace> show recent' output. Example to such output:
// +------------------+-------+--------------+----------------------+--------------------------------------+----------+---------------------+---------------------+------------------+
// |      Tablet      | shard | mysql_schema |     mysql_table      |            migration_uuid            | strategy |  started_timestamp  | completed_timestamp | migration_status |
// +------------------+-------+--------------+----------------------+--------------------------------------+----------+---------------------+---------------------+------------------+
// | zone1-0000003880 |     0 | vt_ks        | vt_onlineddl_test_03 | a0638f6b_ec7b_11ea_9bf8_000d3a9b8a9a | gh-ost   | 2020-09-01 17:50:40 | 2020-09-01 17:50:41 | complete         |
// | zone1-0000003884 |     1 | vt_ks        | vt_onlineddl_test_03 | a0638f6b_ec7b_11ea_9bf8_000d3a9b8a9a | gh-ost   | 2020-09-01 17:50:40 | 2020-09-01 17:50:41 | complete         |
// +------------------+-------+--------------+----------------------+--------------------------------------+----------+---------------------+---------------------+------------------+

func checkRecentMigrations(t *testing.T, uuid string, expectStatusRegexp *regexp.Regexp) {
	result, err := clusterInstance.VtctlclientProcess.OnlineDDLShowRecent(keyspaceName)
	assert.NoError(t, err)
	fmt.Println("# 'vtctlclient OnlineDDL show recent' output (for debug purposes):")
	fmt.Println(result)
	assert.Equal(t, len(clusterInstance.Keyspaces[0].Shards), strings.Count(result, uuid))
	// The word "complete" appears in the column `completed_timestamp`. So we use a regexp to
	// ensure we match exact full word
	m := expectStatusRegexp.FindAllString(result, -1)
	assert.Equal(t, len(clusterInstance.Keyspaces[0].Shards), len(m))
}

// checkCancelMigration attempts to cancel a migration, and expects rejection
func checkCancelMigration(t *testing.T, uuid string) {
	result, err := clusterInstance.VtctlclientProcess.OnlineDDLCancelMigration(keyspaceName, uuid)
	assert.NoError(t, err)
	// The migration has either been complete or failed. We can't cancel it. Expect "zero" response from all tablets
	m := regexp.MustCompile("\b0\b").FindAllString(result, -1)
	assert.Equal(t, len(clusterInstance.Keyspaces[0].Shards), len(m))
}

// checkRetryMigration attempts to retry a migration, and expects rejection
func checkRetryMigration(t *testing.T, uuid string, expectSuccess bool) {
	result, err := clusterInstance.VtctlclientProcess.OnlineDDLRetryMigration(keyspaceName, uuid)
	fmt.Println("# 'vtctlclient OnlineDDL retry <uuid>' output (for debug purposes):")
	fmt.Println(result)
	assert.NoError(t, err)
	// The migration has either been complete or failed. We can't cancel it. Expect "zero" response from all tablets
	var r *regexp.Regexp
	if expectSuccess {
		r = regexp.MustCompile("\b1\b")
	} else {
		r = regexp.MustCompile("\b0\b")
	}
	m := r.FindAllString(result, -1)
	assert.Equal(t, len(clusterInstance.Keyspaces[0].Shards), len(m))
}

// checkMigratedTables checks the CREATE STATEMENT of a table after migration
func checkMigratedTable(t *testing.T, tableName string) {
	expect := "ghost_col"
	for i := range clusterInstance.Keyspaces[0].Shards {
		createStatement := getCreateTableStatement(t, clusterInstance.Keyspaces[0].Shards[i].Vttablets[0], tableName)
		assert.True(t, strings.Contains(createStatement, expect))
	}
}

// getCreateTableStatement returns the CREATE TABLE statement for a given table
func getCreateTableStatement(t *testing.T, tablet *cluster.Vttablet, tableName string) (statement string) {
	queryResult, err := tablet.VttabletProcess.QueryTablet(fmt.Sprintf("show create table %s;", tableName), keyspaceName, true)
	require.Nil(t, err)

	assert.Equal(t, len(queryResult.Rows), 1)
	assert.Equal(t, len(queryResult.Rows[0]), 2) // table name, create statement
	statement = queryResult.Rows[0][1].ToString()
	return statement
}
<|endoftext|>