<commit_before>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package libkbfs

import (
	"fmt"

	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/kbfs/kbfscrypto"

	"github.com/keybase/kbfs/tlf"

	"golang.org/x/net/context"
)

// journalMDOps is an implementation of MDOps that delegates to a
// TLF's mdJournal, if one exists. Specifically, it intercepts put
// calls to write to the journal instead of the MDServer, where
// something else is presumably flushing the journal to the MDServer.
//
// It then intercepts get calls to provide a combined view of the MDs
// from the journal and the server when the journal is
// non-empty. Specifically, if rev is the earliest revision in the
// journal, and BID is the branch ID of the journal (which can only
// have one), then any requests for revisions >= rev on BID will be
// served from the journal instead of the server. If BID is empty,
// i.e. the journal is holding merged revisions, then this means that
// all merged revisions on the server from rev are hidden.
//
// TODO: This makes server updates meaningless for revisions >=
// rev. Fix this.
type journalMDOps struct {
	MDOps
	jServer *JournalServer
}

var _ MDOps = journalMDOps{}

// convertImmutableBareRMDToIRMD decrypts the bare MD into a
// full-fledged RMD.
func (j journalMDOps) convertImmutableBareRMDToIRMD(ctx context.Context,
	ibrmd ImmutableBareRootMetadata, handle *TlfHandle,
	uid keybase1.UID, key kbfscrypto.VerifyingKey) (
	ImmutableRootMetadata, error) {
	// TODO: Avoid having to do this type assertion.
	brmd, ok := ibrmd.BareRootMetadata.(MutableBareRootMetadata)
	if !ok {
		return ImmutableRootMetadata{}, MutableBareRootMetadataNoImplError{}
	}

	rmd := makeRootMetadata(brmd, ibrmd.extra, handle)

	config := j.jServer.config
	pmd, err := decryptMDPrivateData(ctx, config.Codec(), config.Crypto(),
		config.BlockCache(), config.BlockOps(), config.KeyManager(),
		uid, rmd.GetSerializedPrivateMetadata(), rmd, rmd)
	if err != nil {
		return ImmutableRootMetadata{}, err
	}

	rmd.data = pmd
	irmd := MakeImmutableRootMetadata(
		rmd, key, ibrmd.mdID, ibrmd.localTimestamp)
	return irmd, nil
}

// getHeadFromJournal returns the head RootMetadata for the TLF with
// the given ID stored in the journal, assuming it exists and matches
// the given branch ID and merge status. As a special case, if bid is
// NullBranchID and mStatus is Unmerged, the branch ID check is
// skipped.
func (j journalMDOps) getHeadFromJournal(
	ctx context.Context, id tlf.ID, bid BranchID, mStatus MergeStatus,
	handle *TlfHandle) (
	ImmutableRootMetadata, error) {
	tlfJournal, ok := j.jServer.getTLFJournal(id)
	if !ok {
		return ImmutableRootMetadata{}, nil
	}

	head, err := tlfJournal.getMDHead(ctx)
	if err == errTLFJournalDisabled {
		return ImmutableRootMetadata{}, nil
	} else if err != nil {
		return ImmutableRootMetadata{}, err
	}

	if head == (ImmutableBareRootMetadata{}) {
		return ImmutableRootMetadata{}, nil
	}

	if head.MergedStatus() != mStatus {
		return ImmutableRootMetadata{}, nil
	}

	if mStatus == Unmerged && bid != NullBranchID && bid != head.BID() {
		// The given branch ID doesn't match the one in the
		// journal, which can only be an error.
		return ImmutableRootMetadata{},
			fmt.Errorf("Expected branch ID %s, got %s",
				bid, head.BID())
	}

	headBareHandle, err := head.MakeBareTlfHandleWithExtra()
	if err != nil {
		return ImmutableRootMetadata{}, err
	}

	if handle == nil {
		handle, err = MakeTlfHandle(
			ctx, headBareHandle, j.jServer.config.KBPKI())
		if err != nil {
			return ImmutableRootMetadata{}, err
		}
	} else {
		// Check for mutual handle resolution.
		headHandle, err := MakeTlfHandle(ctx, headBareHandle,
			j.jServer.config.KBPKI())
		if err != nil {
			return ImmutableRootMetadata{}, err
		}

		if err := headHandle.MutuallyResolvesTo(ctx, j.jServer.config.Codec(),
			j.jServer.config.KBPKI(), *handle, head.RevisionNumber(),
			head.TlfID(), j.jServer.log); err != nil {
			return ImmutableRootMetadata{}, err
		}
	}

	irmd, err := j.convertImmutableBareRMDToIRMD(
		ctx, head, handle, tlfJournal.uid, tlfJournal.key)
	if err != nil {
		return ImmutableRootMetadata{}, err
	}

	return irmd, nil
}

func (j journalMDOps) getRangeFromJournal(
	ctx context.Context, id tlf.ID, bid BranchID, mStatus MergeStatus,
	start, stop MetadataRevision) (
	[]ImmutableRootMetadata, error) {
	tlfJournal, ok := j.jServer.getTLFJournal(id)
	if !ok {
		return nil, nil
	}

	ibrmds, err := tlfJournal.getMDRange(ctx, start, stop)
	if err == errTLFJournalDisabled {
		return nil, nil
	} else if err != nil {
		return nil, err
	}

	if len(ibrmds) == 0 {
		return nil, nil
	}

	head := ibrmds[len(ibrmds)-1]

	if head.MergedStatus() != mStatus {
		return nil, nil
	}

	if mStatus == Unmerged && bid != NullBranchID && bid != head.BID() {
		// The given branch ID doesn't match the one in the
		// journal, which can only be an error.
		return nil, fmt.Errorf("Expected branch ID %s, got %s",
			bid, head.BID())
	}

	bareHandle, err := head.MakeBareTlfHandleWithExtra()
	if err != nil {
		return nil, err
	}
	handle, err := MakeTlfHandle(ctx, bareHandle, j.jServer.config.KBPKI())
	if err != nil {
		return nil, err
	}

	irmds := make([]ImmutableRootMetadata, 0, len(ibrmds))

	for _, ibrmd := range ibrmds {
		irmd, err := j.convertImmutableBareRMDToIRMD(
			ctx, ibrmd, handle, tlfJournal.uid, tlfJournal.key)
		if err != nil {
			return nil, err
		}

		irmds = append(irmds, irmd)
	}

	return irmds, nil
}

func (j journalMDOps) GetForHandle(
	ctx context.Context, handle *TlfHandle, mStatus MergeStatus) (
	tlf.ID, ImmutableRootMetadata, error) {
	// Need to always consult the server to get the tlfID. No need
	// to optimize this, since all subsequent lookups will be by
	// TLF. Although if we did want to, we could store a handle ->
	// TLF ID mapping with the journals.
	tlfID, rmd, err := j.MDOps.GetForHandle(ctx, handle, mStatus)
	if err != nil {
		return tlf.ID{}, ImmutableRootMetadata{}, err
	}

	if rmd != (ImmutableRootMetadata{}) && (rmd.TlfID() != tlfID) {
		return tlf.ID{}, ImmutableRootMetadata{},
			fmt.Errorf("Expected RMD to have TLF ID %s, but got %s",
				tlfID, rmd.TlfID())
	}

	// If the journal has a head, use that.
	irmd, err := j.getHeadFromJournal(
		ctx, tlfID, NullBranchID, mStatus, handle)
	if err != nil {
		return tlf.ID{}, ImmutableRootMetadata{}, err
	}
	if irmd != (ImmutableRootMetadata{}) {
		return tlf.ID{}, irmd, nil
	}

	// Otherwise, use the server's head.
	return tlfID, rmd, nil
}

// TODO: Combine the two GetForTLF functions in MDOps to avoid the
// need for this helper function.
func (j journalMDOps) getForTLF(
	ctx context.Context, id tlf.ID, bid BranchID, mStatus MergeStatus,
	delegateFn func(context.Context, tlf.ID) (ImmutableRootMetadata, error)) (
	ImmutableRootMetadata, error) {
	// If the journal has a head, use that.
	irmd, err := j.getHeadFromJournal(ctx, id, bid, mStatus, nil)
	if err != nil {
		return ImmutableRootMetadata{}, err
	}
	if irmd != (ImmutableRootMetadata{}) {
		return irmd, nil
	}

	// Otherwise, consult the server instead.
	return delegateFn(ctx, id)
}

func (j journalMDOps) GetForTLF(
	ctx context.Context, id tlf.ID) (ImmutableRootMetadata, error) {
	return j.getForTLF(ctx, id, NullBranchID, Merged, j.MDOps.GetForTLF)
}

func (j journalMDOps) GetUnmergedForTLF(
	ctx context.Context, id tlf.ID, bid BranchID) (
	ImmutableRootMetadata, error) {
	delegateFn := func(ctx context.Context, id tlf.ID) (
		ImmutableRootMetadata, error) {
		return j.MDOps.GetUnmergedForTLF(ctx, id, bid)
	}
	return j.getForTLF(ctx, id, bid, Unmerged, delegateFn)
}

// TODO: Combine the two GetRange functions in MDOps to avoid the need
// for this helper function.
func (j journalMDOps) getRange(
	ctx context.Context, id tlf.ID, bid BranchID, mStatus MergeStatus,
	start, stop MetadataRevision,
	delegateFn func(ctx context.Context, id tlf.ID,
		start, stop MetadataRevision) (
		[]ImmutableRootMetadata, error)) (
	[]ImmutableRootMetadata, error) {
	// Grab the range from the journal first.
	jirmds, err := j.getRangeFromJournal(ctx, id, bid, mStatus, start, stop)
	if err != nil {
		return nil, err
	}

	// If it's empty or disabled, just fall back to the server.
	if len(jirmds) == 0 || err == errTLFJournalDisabled {
		return delegateFn(ctx, id, start, stop)
	}

	// If the first revision from the journal is the first
	// revision we asked for, then just return the range from the
	// journal.
	if jirmds[0].Revision() == start {
		return jirmds, nil
	}

	// Otherwise, fetch the rest from the server and prepend them.
	serverStop := jirmds[0].Revision() - 1
	irmds, err := delegateFn(ctx, id, start, serverStop)
	if err != nil {
		return nil, err
	}

	if len(irmds) == 0 {
		return jirmds, nil
	}

	lastRev := irmds[len(irmds)-1].Revision()
	if lastRev != serverStop {
		return nil, fmt.Errorf(
			"Expected last server rev %d, got %d",
			serverStop, lastRev)
	}

	return append(irmds, jirmds...), nil
}

func (j journalMDOps) GetRange(
	ctx context.Context, id tlf.ID, start, stop MetadataRevision) (
	[]ImmutableRootMetadata, error) {
	return j.getRange(ctx, id, NullBranchID, Merged, start, stop,
		j.MDOps.GetRange)
}

func (j journalMDOps) GetUnmergedRange(
	ctx context.Context, id tlf.ID, bid BranchID,
	start, stop MetadataRevision) ([]ImmutableRootMetadata, error) {
	delegateFn := func(ctx context.Context, id tlf.ID,
		start, stop MetadataRevision) (
		[]ImmutableRootMetadata, error) {
		return j.MDOps.GetUnmergedRange(ctx, id, bid, start, stop)
	}
	return j.getRange(ctx, id, bid, Unmerged, start, stop,
		delegateFn)
}

func (j journalMDOps) Put(ctx context.Context, rmd *RootMetadata) (
	MdID, error) {
	if tlfJournal, ok := j.jServer.getTLFJournal(rmd.TlfID()); ok {
		// Just route to the journal.
		mdID, err := tlfJournal.putMD(ctx, rmd)
		if err != errTLFJournalDisabled {
			return mdID, err
		}
	}

	return j.MDOps.Put(ctx, rmd)
}

func (j journalMDOps) PutUnmerged(ctx context.Context, rmd *RootMetadata) (
	MdID, error) {
	if tlfJournal, ok := j.jServer.getTLFJournal(rmd.TlfID()); ok {
		rmd.SetUnmerged()
		mdID, err := tlfJournal.putMD(ctx, rmd)
		if err != errTLFJournalDisabled {
			return mdID, err
		}
	}

	return j.MDOps.PutUnmerged(ctx, rmd)
}

func (j journalMDOps) PruneBranch(
	ctx context.Context, id tlf.ID, bid BranchID) error {
	if tlfJournal, ok := j.jServer.getTLFJournal(id); ok {
		// Prune the journal, too.
		err := tlfJournal.clearMDs(ctx, bid)
		if err != nil && err != errTLFJournalDisabled {
			return err
		}
	}

	return j.MDOps.PruneBranch(ctx, id, bid)
}

func (j journalMDOps) ResolveBranch(
	ctx context.Context, id tlf.ID, bid BranchID,
	blocksToDelete []BlockID, rmd *RootMetadata) (MdID, error) {
	if tlfJournal, ok := j.jServer.getTLFJournal(id); ok {
		mdID, err := tlfJournal.resolveBranch(
			ctx, bid, blocksToDelete, rmd, rmd.extra)
		if err != errTLFJournalDisabled {
			return mdID, err
		}
	}

	return j.MDOps.ResolveBranch(ctx, id, bid, blocksToDelete, rmd)
}
<commit_msg>journal_md_ops: always use Merged mStatus to find TLF ID<commit_after>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package libkbfs

import (
	"fmt"

	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/kbfs/kbfscrypto"

	"github.com/keybase/kbfs/tlf"

	"golang.org/x/net/context"
)

// journalMDOps is an implementation of MDOps that delegates to a
// TLF's mdJournal, if one exists. Specifically, it intercepts put
// calls to write to the journal instead of the MDServer, where
// something else is presumably flushing the journal to the MDServer.
//
// It then intercepts get calls to provide a combined view of the MDs
// from the journal and the server when the journal is
// non-empty. Specifically, if rev is the earliest revision in the
// journal, and BID is the branch ID of the journal (which can only
// have one), then any requests for revisions >= rev on BID will be
// served from the journal instead of the server. If BID is empty,
// i.e. the journal is holding merged revisions, then this means that
// all merged revisions on the server from rev are hidden.
//
// TODO: This makes server updates meaningless for revisions >=
// rev. Fix this.
type journalMDOps struct {
	MDOps
	jServer *JournalServer
}

var _ MDOps = journalMDOps{}

// convertImmutableBareRMDToIRMD decrypts the bare MD into a
// full-fledged RMD.
func (j journalMDOps) convertImmutableBareRMDToIRMD(ctx context.Context,
	ibrmd ImmutableBareRootMetadata, handle *TlfHandle,
	uid keybase1.UID, key kbfscrypto.VerifyingKey) (
	ImmutableRootMetadata, error) {
	// TODO: Avoid having to do this type assertion.
	brmd, ok := ibrmd.BareRootMetadata.(MutableBareRootMetadata)
	if !ok {
		return ImmutableRootMetadata{}, MutableBareRootMetadataNoImplError{}
	}

	rmd := makeRootMetadata(brmd, ibrmd.extra, handle)

	config := j.jServer.config
	pmd, err := decryptMDPrivateData(ctx, config.Codec(), config.Crypto(),
		config.BlockCache(), config.BlockOps(), config.KeyManager(),
		uid, rmd.GetSerializedPrivateMetadata(), rmd, rmd)
	if err != nil {
		return ImmutableRootMetadata{}, err
	}

	rmd.data = pmd
	irmd := MakeImmutableRootMetadata(
		rmd, key, ibrmd.mdID, ibrmd.localTimestamp)
	return irmd, nil
}

// getHeadFromJournal returns the head RootMetadata for the TLF with
// the given ID stored in the journal, assuming it exists and matches
// the given branch ID and merge status. As a special case, if bid is
// NullBranchID and mStatus is Unmerged, the branch ID check is
// skipped.
func (j journalMDOps) getHeadFromJournal(
	ctx context.Context, id tlf.ID, bid BranchID, mStatus MergeStatus,
	handle *TlfHandle) (
	ImmutableRootMetadata, error) {
	tlfJournal, ok := j.jServer.getTLFJournal(id)
	if !ok {
		return ImmutableRootMetadata{}, nil
	}

	head, err := tlfJournal.getMDHead(ctx)
	if err == errTLFJournalDisabled {
		return ImmutableRootMetadata{}, nil
	} else if err != nil {
		return ImmutableRootMetadata{}, err
	}

	if head == (ImmutableBareRootMetadata{}) {
		return ImmutableRootMetadata{}, nil
	}

	if head.MergedStatus() != mStatus {
		return ImmutableRootMetadata{}, nil
	}

	if mStatus == Unmerged && bid != NullBranchID && bid != head.BID() {
		// The given branch ID doesn't match the one in the
		// journal, which can only be an error.
		return ImmutableRootMetadata{},
			fmt.Errorf("Expected branch ID %s, got %s",
				bid, head.BID())
	}

	headBareHandle, err := head.MakeBareTlfHandleWithExtra()
	if err != nil {
		return ImmutableRootMetadata{}, err
	}

	if handle == nil {
		handle, err = MakeTlfHandle(
			ctx, headBareHandle, j.jServer.config.KBPKI())
		if err != nil {
			return ImmutableRootMetadata{}, err
		}
	} else {
		// Check for mutual handle resolution.
		headHandle, err := MakeTlfHandle(ctx, headBareHandle,
			j.jServer.config.KBPKI())
		if err != nil {
			return ImmutableRootMetadata{}, err
		}

		if err := headHandle.MutuallyResolvesTo(ctx, j.jServer.config.Codec(),
			j.jServer.config.KBPKI(), *handle, head.RevisionNumber(),
			head.TlfID(), j.jServer.log); err != nil {
			return ImmutableRootMetadata{}, err
		}
	}

	irmd, err := j.convertImmutableBareRMDToIRMD(
		ctx, head, handle, tlfJournal.uid, tlfJournal.key)
	if err != nil {
		return ImmutableRootMetadata{}, err
	}

	return irmd, nil
}

func (j journalMDOps) getRangeFromJournal(
	ctx context.Context, id tlf.ID, bid BranchID, mStatus MergeStatus,
	start, stop MetadataRevision) (
	[]ImmutableRootMetadata, error) {
	tlfJournal, ok := j.jServer.getTLFJournal(id)
	if !ok {
		return nil, nil
	}

	ibrmds, err := tlfJournal.getMDRange(ctx, start, stop)
	if err == errTLFJournalDisabled {
		return nil, nil
	} else if err != nil {
		return nil, err
	}

	if len(ibrmds) == 0 {
		return nil, nil
	}

	head := ibrmds[len(ibrmds)-1]

	if head.MergedStatus() != mStatus {
		return nil, nil
	}

	if mStatus == Unmerged && bid != NullBranchID && bid != head.BID() {
		// The given branch ID doesn't match the one in the
		// journal, which can only be an error.
		return nil, fmt.Errorf("Expected branch ID %s, got %s",
			bid, head.BID())
	}

	bareHandle, err := head.MakeBareTlfHandleWithExtra()
	if err != nil {
		return nil, err
	}
	handle, err := MakeTlfHandle(ctx, bareHandle, j.jServer.config.KBPKI())
	if err != nil {
		return nil, err
	}

	irmds := make([]ImmutableRootMetadata, 0, len(ibrmds))

	for _, ibrmd := range ibrmds {
		irmd, err := j.convertImmutableBareRMDToIRMD(
			ctx, ibrmd, handle, tlfJournal.uid, tlfJournal.key)
		if err != nil {
			return nil, err
		}

		irmds = append(irmds, irmd)
	}

	return irmds, nil
}

func (j journalMDOps) GetForHandle(
	ctx context.Context, handle *TlfHandle, mStatus MergeStatus) (
	tlf.ID, ImmutableRootMetadata, error) {
	// Need to always consult the server to get the tlfID. No need to
	// optimize this, since all subsequent lookups will be by
	// TLF. Although if we did want to, we could store a handle -> TLF
	// ID mapping with the journals.  If we are looking for an
	// unmerged head, that exists only in the journal, so check the
	// remote server only to get the TLF ID.
	remoteMStatus := mStatus
	if mStatus == Unmerged {
		remoteMStatus = Merged
	}
	tlfID, rmd, err := j.MDOps.GetForHandle(ctx, handle, remoteMStatus)
	if err != nil {
		return tlf.ID{}, ImmutableRootMetadata{}, err
	}

	if rmd != (ImmutableRootMetadata{}) && (rmd.TlfID() != tlfID) {
		return tlf.ID{}, ImmutableRootMetadata{},
			fmt.Errorf("Expected RMD to have TLF ID %s, but got %s",
				tlfID, rmd.TlfID())
	}

	// If the journal has a head, use that.
	irmd, err := j.getHeadFromJournal(
		ctx, tlfID, NullBranchID, mStatus, handle)
	if err != nil {
		return tlf.ID{}, ImmutableRootMetadata{}, err
	}
	if irmd != (ImmutableRootMetadata{}) {
		return tlf.ID{}, irmd, nil
	}
	if remoteMStatus != mStatus {
		return tlfID, ImmutableRootMetadata{}, nil
	}

	// Otherwise, use the server's head.
	return tlfID, rmd, nil
}

// TODO: Combine the two GetForTLF functions in MDOps to avoid the
// need for this helper function.
func (j journalMDOps) getForTLF(
	ctx context.Context, id tlf.ID, bid BranchID, mStatus MergeStatus,
	delegateFn func(context.Context, tlf.ID) (ImmutableRootMetadata, error)) (
	ImmutableRootMetadata, error) {
	// If the journal has a head, use that.
	irmd, err := j.getHeadFromJournal(ctx, id, bid, mStatus, nil)
	if err != nil {
		return ImmutableRootMetadata{}, err
	}
	if irmd != (ImmutableRootMetadata{}) {
		return irmd, nil
	}

	// Otherwise, consult the server instead.
	return delegateFn(ctx, id)
}

func (j journalMDOps) GetForTLF(
	ctx context.Context, id tlf.ID) (ImmutableRootMetadata, error) {
	return j.getForTLF(ctx, id, NullBranchID, Merged, j.MDOps.GetForTLF)
}

func (j journalMDOps) GetUnmergedForTLF(
	ctx context.Context, id tlf.ID, bid BranchID) (
	ImmutableRootMetadata, error) {
	delegateFn := func(ctx context.Context, id tlf.ID) (
		ImmutableRootMetadata, error) {
		return j.MDOps.GetUnmergedForTLF(ctx, id, bid)
	}
	return j.getForTLF(ctx, id, bid, Unmerged, delegateFn)
}

// TODO: Combine the two GetRange functions in MDOps to avoid the need
// for this helper function.
func (j journalMDOps) getRange(
	ctx context.Context, id tlf.ID, bid BranchID, mStatus MergeStatus,
	start, stop MetadataRevision,
	delegateFn func(ctx context.Context, id tlf.ID,
		start, stop MetadataRevision) (
		[]ImmutableRootMetadata, error)) (
	[]ImmutableRootMetadata, error) {
	// Grab the range from the journal first.
	jirmds, err := j.getRangeFromJournal(ctx, id, bid, mStatus, start, stop)
	if err != nil {
		return nil, err
	}

	// If it's empty or disabled, just fall back to the server.
	if len(jirmds) == 0 || err == errTLFJournalDisabled {
		return delegateFn(ctx, id, start, stop)
	}

	// If the first revision from the journal is the first
	// revision we asked for, then just return the range from the
	// journal.
	if jirmds[0].Revision() == start {
		return jirmds, nil
	}

	// Otherwise, fetch the rest from the server and prepend them.
	serverStop := jirmds[0].Revision() - 1
	irmds, err := delegateFn(ctx, id, start, serverStop)
	if err != nil {
		return nil, err
	}

	if len(irmds) == 0 {
		return jirmds, nil
	}

	lastRev := irmds[len(irmds)-1].Revision()
	if lastRev != serverStop {
		return nil, fmt.Errorf(
			"Expected last server rev %d, got %d",
			serverStop, lastRev)
	}

	return append(irmds, jirmds...), nil
}

func (j journalMDOps) GetRange(
	ctx context.Context, id tlf.ID, start, stop MetadataRevision) (
	[]ImmutableRootMetadata, error) {
	return j.getRange(ctx, id, NullBranchID, Merged, start, stop,
		j.MDOps.GetRange)
}

func (j journalMDOps) GetUnmergedRange(
	ctx context.Context, id tlf.ID, bid BranchID,
	start, stop MetadataRevision) ([]ImmutableRootMetadata, error) {
	delegateFn := func(ctx context.Context, id tlf.ID,
		start, stop MetadataRevision) (
		[]ImmutableRootMetadata, error) {
		return j.MDOps.GetUnmergedRange(ctx, id, bid, start, stop)
	}
	return j.getRange(ctx, id, bid, Unmerged, start, stop,
		delegateFn)
}

func (j journalMDOps) Put(ctx context.Context, rmd *RootMetadata) (
	MdID, error) {
	if tlfJournal, ok := j.jServer.getTLFJournal(rmd.TlfID()); ok {
		// Just route to the journal.
		mdID, err := tlfJournal.putMD(ctx, rmd)
		if err != errTLFJournalDisabled {
			return mdID, err
		}
	}

	return j.MDOps.Put(ctx, rmd)
}

func (j journalMDOps) PutUnmerged(ctx context.Context, rmd *RootMetadata) (
	MdID, error) {
	if tlfJournal, ok := j.jServer.getTLFJournal(rmd.TlfID()); ok {
		rmd.SetUnmerged()
		mdID, err := tlfJournal.putMD(ctx, rmd)
		if err != errTLFJournalDisabled {
			return mdID, err
		}
	}

	return j.MDOps.PutUnmerged(ctx, rmd)
}

func (j journalMDOps) PruneBranch(
	ctx context.Context, id tlf.ID, bid BranchID) error {
	if tlfJournal, ok := j.jServer.getTLFJournal(id); ok {
		// Prune the journal, too.
		err := tlfJournal.clearMDs(ctx, bid)
		if err != nil && err != errTLFJournalDisabled {
			return err
		}
	}

	return j.MDOps.PruneBranch(ctx, id, bid)
}

func (j journalMDOps) ResolveBranch(
	ctx context.Context, id tlf.ID, bid BranchID,
	blocksToDelete []BlockID, rmd *RootMetadata) (MdID, error) {
	if tlfJournal, ok := j.jServer.getTLFJournal(id); ok {
		mdID, err := tlfJournal.resolveBranch(
			ctx, bid, blocksToDelete, rmd, rmd.extra)
		if err != errTLFJournalDisabled {
			return mdID, err
		}
	}

	return j.MDOps.ResolveBranch(ctx, id, bid, blocksToDelete, rmd)
}
<|endoftext|>


<commit_before>// Command protoc-gen-doc is a Protobuf plugin for documentation generation.
//
// More information about Protobuf can be found at:
//
// 	https://developers.google.com/protocol-buffers/
//
package main // import "sourcegraph.com/sourcegraph/prototools/cmd/protoc-gen-doc"

import (
	"go/build"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"text/template"

	"github.com/golang/protobuf/proto"
	"sourcegraph.com/sourcegraph/prototools/tmpl"
)

// PathDir returns the absolute path to a file given a relative one in one of
// the $GOPATH directories. If it cannot be resolved, relPath itself is
// returned.
func PathDir(relPath string) string {
	// Test again each directory listed in $GOPATH
	for _, path := range filepath.SplitList(build.Default.GOPATH) {
		path = filepath.Join(path, relPath)
		if _, err := os.Stat(path); err == nil {
			return path
		}
	}
	return relPath
}

func main() {
	// Configure logging.
	log.SetFlags(0)
	log.SetPrefix("protoc-gen-doc: ")

	// Create a template generator.
	g := tmpl.New()

	// Read input from the protoc compiler.
	data, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		log.Fatal(err, ": failed to read input")
	}

	// Unmarshal the protoc generation request.
	if err := proto.Unmarshal(data, g.Request); err != nil {
		log.Fatal(err, ": failed to parse input proto")
	}
	if len(g.Request.FileToGenerate) == 0 {
		log.Fatal(err, ": no input files")
	}

	// Parse the command-line parameters and determine the correct template file
	// to execute.
	g.ParseParams()
	var tmplPath string
	if v, ok := g.Param["template"]; ok {
		tmplPath = v
	} else {
		tmplPath = PathDir("src/sourcegraph.com/sourcegraph/prototools/templates/tmpl.html")
	}

	// Load up the template, preloading the function map (or else the functions
	// will fail when called).
	t := template.New("").Funcs(tmpl.Preload)
	g.Template, err = t.ParseGlob(tmplPath)
	if err != nil {
		log.Fatal(err, ": failed to parse templates")
	}
	g.Template = g.Template.Templates()[0]

	// Perform generation.
	response, err := g.Generate()
	if err != nil {
		log.Fatal(err, ": failed to generate")
	}

	// Marshal the results and write back to the protoc compiler.
	data, err = proto.Marshal(response)
	if err != nil {
		log.Fatal(err, ": failed to marshal output proto")
	}
	_, err = os.Stdout.Write(data)
	if err != nil {
		log.Fatal(err, ": failed to write output proto")
	}
}
<commit_msg>cmd/protoc-gen-doc: fix incorrect template selection and half-writes<commit_after>// Command protoc-gen-doc is a Protobuf plugin for documentation generation.
//
// More information about Protobuf can be found at:
//
// 	https://developers.google.com/protocol-buffers/
//
package main // import "sourcegraph.com/sourcegraph/prototools/cmd/protoc-gen-doc"

import (
	"bytes"
	"go/build"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"text/template"

	"github.com/golang/protobuf/proto"
	"sourcegraph.com/sourcegraph/prototools/tmpl"
)

// PathDir returns the absolute path to a file given a relative one in one of
// the $GOPATH directories. If it cannot be resolved, relPath itself is
// returned.
func PathDir(relPath string) string {
	// Test again each directory listed in $GOPATH
	for _, path := range filepath.SplitList(build.Default.GOPATH) {
		path = filepath.Join(path, relPath)
		if _, err := os.Stat(path); err == nil {
			return path
		}
	}
	return relPath
}

func main() {
	// Configure logging.
	log.SetFlags(0)
	log.SetPrefix("protoc-gen-doc: ")

	// Create a template generator.
	g := tmpl.New()

	// Read input from the protoc compiler.
	data, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		log.Fatal(err, ": failed to read input")
	}

	// Unmarshal the protoc generation request.
	if err := proto.Unmarshal(data, g.Request); err != nil {
		log.Fatal(err, ": failed to parse input proto")
	}
	if len(g.Request.FileToGenerate) == 0 {
		log.Fatal(err, ": no input files")
	}

	// Parse the command-line parameters and determine the correct template file
	// to execute.
	g.ParseParams()
	var tmplPath string
	if v, ok := g.Param["template"]; ok {
		tmplPath = v
	} else {
		tmplPath = PathDir("src/sourcegraph.com/sourcegraph/prototools/templates/tmpl.html")
	}

	// Load up the template, preloading the function map (or else the functions
	// will fail when called).
	t := template.New("").Funcs(tmpl.Preload)
	g.Template, err = t.ParseGlob(tmplPath)
	if err != nil {
		log.Fatal(err, ": failed to parse templates")
	}
	g.Template = g.Template.Lookup(filepath.Base(tmplPath))

	// Perform generation.
	response, err := g.Generate()
	if err != nil {
		log.Fatal(err, ": failed to generate")
	}

	// Marshal the results and write back to the protoc compiler.
	data, err = proto.Marshal(response)
	if err != nil {
		log.Fatal(err, ": failed to marshal output proto")
	}
	_, err = io.Copy(os.Stdout, bytes.NewReader(data))
	if err != nil {
		log.Fatal(err, ": failed to write output proto")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"flag"
	"net/http"

	assetfs "github.com/elazarl/go-bindata-assetfs"
	"github.com/mmikulicic/stringlist"
	"within.website/ln"
	"within.website/ln/opname"
	"within.website/x/internal"
	"within.website/x/vanity"
)

//go:generate go-bindata -pkg main static

var (
	domain         = flag.String("domain", "within.website", "domain this is run on")
	githubUsername = flag.String("github-user", "Xe", "GitHub username for GitHub repos")
	gogsDomain     = flag.String("gogs-url", "https://git.xeserv.us", "Gogs domain to use")
	gogsUsername   = flag.String("gogs-username", "xena", "Gogs username for above Gogs instance")
	port           = flag.String("port", "2134", "HTTP port to listen on")
	goProxyServer  = flag.String("go-proxy-server", "https://cache.greedo.xeserv.us", "go proxy server to point to for go module clients")

	githubRepos = stringlist.Flag("github-repo", "list of GitHub repositories to use")
	gogsRepos   = stringlist.Flag("gogs-repo", "list of Gogs repositories to use")
)

var githubReposDefault = []string{
	"ln",
	"x",
	"eclier",
	"gluanetrc",
	"xultybau",
	"johaus",
	"confyg",
	"derpigo",
}

var gogsReposDefault = []string{
	"gorqlite",
}

func main() {
	internal.HandleStartup()
	ctx := opname.With(context.Background(), "main")
	ctx = ln.WithF(ctx, ln.F{
		"domain":       *domain,
		"proxy_server": *goProxyServer,
	})

	if len(*githubRepos) == 0 {
		*githubRepos = githubReposDefault
	}

	if len(*gogsRepos) == 0 {
		*gogsRepos = gogsReposDefault
	}

	for _, repo := range *githubRepos {
		http.Handle("/"+repo, vanity.GitHubHandler(*domain+"/"+repo, *githubUsername, repo, "https"))
		http.Handle("/"+repo+"/", vanity.GitHubHandler(*domain+"/"+repo, *githubUsername, repo, "https"))

		ln.Log(ctx, ln.F{"github_repo": repo, "github_user": *githubUsername}, ln.Info("adding github repo"))
	}

	for _, repo := range *gogsRepos {
		http.Handle("/"+repo, vanity.GogsHandler(*domain+"/"+repo, *gogsDomain, *gogsUsername, repo, "https"))
		http.Handle("/"+repo+"/", vanity.GogsHandler(*domain+"/"+repo, *gogsDomain, *gogsUsername, repo, "https"))

		ln.Log(ctx, ln.F{"gogs_domain": *gogsDomain, "gogs_username": *gogsUsername, "gogs_repo": repo}, ln.Info("adding gogs repo"))
	}

	http.Handle("/static/", http.FileServer(
		&assetfs.AssetFS{
			Asset:    Asset,
			AssetDir: AssetDir,
		},
	))

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/" {
			http.NotFound(w, r)
			return
		}

		w.Header().Add("Content-Type", "text/html")
		w.Write([]byte(indexTemplate))
	})

	http.HandleFunc("/.x.botinfo", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Add("Content-Type", "text/html")
		w.Write([]byte(botInfoPage))
	})

	ln.Log(ctx, ln.F{"port": *port}, ln.Info("Listening on HTTP"))
	http.ListenAndServe(":"+*port, nil)

}

const indexTemplate = `<!DOCTYPE html>

		<title>within.website Go Packages</title>
		<link rel="stylesheet" href="/static/gruvbox.css">
		<meta name="viewport" content="width=device-width, initial-scale=1.0" />

within.website Go Packages

confyg - A generic configuration file parser based on the go modfile parser

derpigo - A simple wrapper to the Derpibooru API

eclier - A go+lua command line application framework

gluanetrc - A GopherLua binding for github.com/dickeyxxx/netrc

gorqlite - A driver for rqlite

johaus - Lojban parsing

ln - Key->value based logging made context-aware and simple

x - Experiments, toys and tinkering (many subpackages)

Need help with these packages? Inquire Within.

`

const botInfoPage = `<link rel="stylesheet" href="/static/gruvbox.css">
<title>x repo bots</title>

x repo bots

Hello, if you are reading this, you have found this URL in your access logs.

If one of these programs is doing something you don't want them to do, please contact me or open an issue here.

`
<commit_msg>cmd/within.website: add olin<commit_after>package main

import (
	"context"
	"flag"
	"net/http"

	assetfs "github.com/elazarl/go-bindata-assetfs"
	"github.com/mmikulicic/stringlist"
	"within.website/ln"
	"within.website/ln/opname"
	"within.website/x/internal"
	"within.website/x/vanity"
)

//go:generate go-bindata -pkg main static

var (
	domain         = flag.String("domain", "within.website", "domain this is run on")
	githubUsername = flag.String("github-user", "Xe", "GitHub username for GitHub repos")
	gogsDomain     = flag.String("gogs-url", "https://git.xeserv.us", "Gogs domain to use")
	gogsUsername   = flag.String("gogs-username", "xena", "Gogs username for above Gogs instance")
	port           = flag.String("port", "2134", "HTTP port to listen on")
	goProxyServer  = flag.String("go-proxy-server", "https://cache.greedo.xeserv.us", "go proxy server to point to for go module clients")

	githubRepos = stringlist.Flag("github-repo", "list of GitHub repositories to use")
	gogsRepos   = stringlist.Flag("gogs-repo", "list of Gogs repositories to use")
)

var githubReposDefault = []string{
	"ln",
	"x",
	"eclier",
	"gluanetrc",
	"xultybau",
	"johaus",
	"confyg",
	"derpigo",
	"olin",
}

var gogsReposDefault = []string{
	"gorqlite",
}

func main() {
	internal.HandleStartup()
	ctx := opname.With(context.Background(), "main")
	ctx = ln.WithF(ctx, ln.F{
		"domain":       *domain,
		"proxy_server": *goProxyServer,
	})

	if len(*githubRepos) == 0 {
		*githubRepos = githubReposDefault
	}

	if len(*gogsRepos) == 0 {
		*gogsRepos = gogsReposDefault
	}

	for _, repo := range *githubRepos {
		http.Handle("/"+repo, vanity.GitHubHandler(*domain+"/"+repo, *githubUsername, repo, "https"))
		http.Handle("/"+repo+"/", vanity.GitHubHandler(*domain+"/"+repo, *githubUsername, repo, "https"))

		ln.Log(ctx, ln.F{"github_repo": repo, "github_user": *githubUsername}, ln.Info("adding github repo"))
	}

	for _, repo := range *gogsRepos {
		http.Handle("/"+repo, vanity.GogsHandler(*domain+"/"+repo, *gogsDomain, *gogsUsername, repo, "https"))
		http.Handle("/"+repo+"/", vanity.GogsHandler(*domain+"/"+repo, *gogsDomain, *gogsUsername, repo, "https"))

		ln.Log(ctx, ln.F{"gogs_domain": *gogsDomain, "gogs_username": *gogsUsername, "gogs_repo": repo}, ln.Info("adding gogs repo"))
	}

	http.Handle("/static/", http.FileServer(
		&assetfs.AssetFS{
			Asset:    Asset,
			AssetDir: AssetDir,
		},
	))

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/" {
			http.NotFound(w, r)
			return
		}

		w.Header().Add("Content-Type", "text/html")
		w.Write([]byte(indexTemplate))
	})

	http.HandleFunc("/.x.botinfo", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Add("Content-Type", "text/html")
		w.Write([]byte(botInfoPage))
	})

	ln.Log(ctx, ln.F{"port": *port}, ln.Info("Listening on HTTP"))
	http.ListenAndServe(":"+*port, nil)

}

const indexTemplate = `<!DOCTYPE html>

		<title>within.website Go Packages</title>
		<link rel="stylesheet" href="/static/gruvbox.css">
		<meta name="viewport" content="width=device-width, initial-scale=1.0" />

within.website Go Packages

confyg - A generic configuration file parser based on the go modfile parser

derpigo - A simple wrapper to the Derpibooru API

eclier - A go+lua command line application framework

gluanetrc - A GopherLua binding for github.com/dickeyxxx/netrc

gorqlite - A driver for rqlite

johaus - Lojban parsing

ln - Key->value based logging made context-aware and simple

olin - WebAssembly on the server

x - Experiments, toys and tinkering (many subpackages)

Need help with these packages? Inquire Within.

`

const botInfoPage = `<link rel="stylesheet" href="/static/gruvbox.css">
<title>x repo bots</title>

x repo bots

Hello, if you are reading this, you have found this URL in your access logs.

If one of these programs is doing something you don't want them to do, please contact me or open an issue here.

`
<|endoftext|>


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

// Package fasthttprouter is a trie based high performance HTTP request router.
//
// A trivial example is:
//
// package main

// import (
//     "fmt"
//     "log"
//
//     "github.com/buaazp/fasthttprouter"
//     "github.com/valyala/fasthttp"
// )

// func Index(ctx *fasthttp.RequestCtx) {
//     fmt.Fprint(ctx, "Welcome!\n")
// }

// func Hello(ctx *fasthttp.RequestCtx) {
//     fmt.Fprintf(ctx, "hello, %s!\n", ctx.UserValue("name"))
// }

// func main() {
//     router := fasthttprouter.New()
//     router.GET("/", Index)
//     router.GET("/hello/:name", Hello)

//     log.Fatal(fasthttp.ListenAndServe(":8080", router.Handler))
// }
//
// The router matches incoming requests by the request method and the path.
// If a handle is registered for this path and method, the router delegates the
// request to that function.
// For the methods GET, POST, PUT, PATCH and DELETE shortcut functions exist to
// register handles, for all other methods router.Handle can be used.
//
// The registered path, against which the router matches incoming requests, can
// contain two types of parameters:
//  Syntax    Type
//  :name     named parameter
//  *name     catch-all parameter
//
// Named parameters are dynamic path segments. They match anything until the
// next '/' or the path end:
//  Path: /blog/:category/:post
//
//  Requests:
//   /blog/go/request-routers            match: category="go", post="request-routers"
//   /blog/go/request-routers/           no match, but the router would redirect
//   /blog/go/                           no match
//   /blog/go/request-routers/comments   no match
//
// Catch-all parameters match anything until the path end, including the
// directory index (the '/' before the catch-all). Since they match anything
// until the end, catch-all parameters must always be the final path element.
//  Path: /files/*filepath
//
//  Requests:
//   /files/                             match: filepath="/"
//   /files/LICENSE                      match: filepath="/LICENSE"
//   /files/templates/article.html       match: filepath="/templates/article.html"
//   /files                              no match, but the router would redirect
//
// The value of parameters is inside ctx.UserValue
// To retrieve the value of a parameter:
//  // use the name of the parameter
//  user := ps.UserValue("user")
//

package fasthttprouter

import (
	"strings"

	"github.com/valyala/fasthttp"
)

var (
	defaultContentType = []byte("text/plain; charset=utf-8")
	questionMark       = []byte("?")
)

// Router is a http.Handler which can be used to dispatch requests to different
// handler functions via configurable routes
type Router struct {
	trees map[string]*node

	// Enables automatic redirection if the current route can't be matched but a
	// handler for the path with (without) the trailing slash exists.
	// For example if /foo/ is requested but a route only exists for /foo, the
	// client is redirected to /foo with http status code 301 for GET requests
	// and 307 for all other request methods.
	RedirectTrailingSlash bool

	// If enabled, the router tries to fix the current request path, if no
	// handle is registered for it.
	// First superfluous path elements like ../ or // are removed.
	// Afterwards the router does a case-insensitive lookup of the cleaned path.
	// If a handle can be found for this route, the router makes a redirection
	// to the corrected path with status code 301 for GET requests and 307 for
	// all other request methods.
	// For example /FOO and /..//Foo could be redirected to /foo.
	// RedirectTrailingSlash is independent of this option.
	RedirectFixedPath bool

	// If enabled, the router checks if another method is allowed for the
	// current route, if the current request can not be routed.
	// If this is the case, the request is answered with 'Method Not Allowed'
	// and HTTP status code 405.
	// If no other Method is allowed, the request is delegated to the NotFound
	// handler.
	HandleMethodNotAllowed bool

	// If enabled, the router automatically replies to OPTIONS requests.
	// Custom OPTIONS handlers take priority over automatic replies.
	HandleOPTIONS bool

	// Configurable http.Handler which is called when no matching route is
	// found. If it is not set, http.NotFound is used.
	NotFound fasthttp.RequestHandler

	// Configurable http.Handler which is called when a request
	// cannot be routed and HandleMethodNotAllowed is true.
	// If it is not set, http.Error with http.StatusMethodNotAllowed is used.
	// The "Allow" header with allowed request methods is set before the handler
	// is called.
	MethodNotAllowed fasthttp.RequestHandler

	// Function to handle panics recovered from http handlers.
	// It should be used to generate a error page and return the http error code
	// 500 (Internal Server Error).
	// The handler can be used to keep your server from crashing because of
	// unrecovered panics.
	PanicHandler func(*fasthttp.RequestCtx, interface{})
}

// New returns a new initialized Router.
// Path auto-correction, including trailing slashes, is enabled by default.
func New() *Router {
	return &Router{
		RedirectTrailingSlash:  true,
		RedirectFixedPath:      true,
		HandleMethodNotAllowed: true,
		HandleOPTIONS:          true,
	}
}

// GET is a shortcut for router.Handle("GET", path, handle)
func (r *Router) GET(path string, handle fasthttp.RequestHandler) {
	r.Handle("GET", path, handle)
}

// HEAD is a shortcut for router.Handle("HEAD", path, handle)
func (r *Router) HEAD(path string, handle fasthttp.RequestHandler) {
	r.Handle("HEAD", path, handle)
}

// OPTIONS is a shortcut for router.Handle("OPTIONS", path, handle)
func (r *Router) OPTIONS(path string, handle fasthttp.RequestHandler) {
	r.Handle("OPTIONS", path, handle)
}

// POST is a shortcut for router.Handle("POST", path, handle)
func (r *Router) POST(path string, handle fasthttp.RequestHandler) {
	r.Handle("POST", path, handle)
}

// PUT is a shortcut for router.Handle("PUT", path, handle)
func (r *Router) PUT(path string, handle fasthttp.RequestHandler) {
	r.Handle("PUT", path, handle)
}

// PATCH is a shortcut for router.Handle("PATCH", path, handle)
func (r *Router) PATCH(path string, handle fasthttp.RequestHandler) {
	r.Handle("PATCH", path, handle)
}

// DELETE is a shortcut for router.Handle("DELETE", path, handle)
func (r *Router) DELETE(path string, handle fasthttp.RequestHandler) {
	r.Handle("DELETE", path, handle)
}

// Handle registers a new request handle with the given path and method.
//
// For GET, POST, PUT, PATCH and DELETE requests the respective shortcut
// functions can be used.
//
// This function is intended for bulk loading and to allow the usage of less
// frequently used, non-standardized or custom methods (e.g. for internal
// communication with a proxy).
func (r *Router) Handle(method, path string, handle fasthttp.RequestHandler) {
	if path[0] != '/' {
		panic("path must begin with '/' in path '" + path + "'")
	}

	if r.trees == nil {
		r.trees = make(map[string]*node)
	}

	root := r.trees[method]
	if root == nil {
		root = new(node)
		r.trees[method] = root
	}

	root.addRoute(path, handle)
}

// ServeFiles serves files from the given file system root.
// The path must end with "/*filepath", files are then served from the local
// path /defined/root/dir/*filepath.
// For example if root is "/etc" and *filepath is "passwd", the local file
// "/etc/passwd" would be served.
// Internally a http.FileServer is used, therefore http.NotFound is used instead
// of the Router's NotFound handler.
//     router.ServeFiles("/src/*filepath", "/var/www")
func (r *Router) ServeFiles(path string, rootPath string) {
	if len(path) < 10 || path[len(path)-10:] != "/*filepath" {
		panic("path must end with /*filepath in path '" + path + "'")
	}
	prefix := path[:len(path)-10]

	fileHandler := fasthttp.FSHandler(rootPath, strings.Count(prefix, "/"))

	r.GET(path, func(ctx *fasthttp.RequestCtx) {
		fileHandler(ctx)
	})
}

func (r *Router) recv(ctx *fasthttp.RequestCtx) {
	if rcv := recover(); rcv != nil {
		r.PanicHandler(ctx, rcv)
	}
}

// Lookup allows the manual lookup of a method + path combo.
// This is e.g. useful to build a framework around this router.
// If the path was found, it returns the handle function and the path parameter
// values. Otherwise the third return value indicates whether a redirection to
// the same path with an extra / without the trailing slash should be performed.
func (r *Router) Lookup(method, path string, ctx *fasthttp.RequestCtx) (fasthttp.RequestHandler, bool) {
	if root := r.trees[method]; root != nil {
		return root.getValue(path, ctx)
	}
	return nil, false
}

func (r *Router) allowed(path, reqMethod string) (allow string) {
	if path == "*" || path == "/*" { // server-wide
		for method := range r.trees {
			if method == "OPTIONS" {
				continue
			}

			// add request method to list of allowed methods
			if len(allow) == 0 {
				allow = method
			} else {
				allow += ", " + method
			}
		}
	} else { // specific path
		for method := range r.trees {
			// Skip the requested method - we already tried this one
			if method == reqMethod || method == "OPTIONS" {
				continue
			}

			handle, _ := r.trees[method].getValue(path, nil)
			if handle != nil {
				// add request method to list of allowed methods
				if len(allow) == 0 {
					allow = method
				} else {
					allow += ", " + method
				}
			}
		}
	}
	if len(allow) > 0 {
		allow += ", OPTIONS"
	}
	return
}

// Handler makes the router implement the fasthttp.ListenAndServe interface.
func (r *Router) Handler(ctx *fasthttp.RequestCtx) {
	if r.PanicHandler != nil {
		defer r.recv(ctx)
	}

	path := string(ctx.Path())
	method := string(ctx.Method())
	if root := r.trees[method]; root != nil {
		if f, tsr := root.getValue(path, ctx); f != nil {
			f(ctx)
			return
		} else if method != "CONNECT" && path != "/" {
			code := 301 // Permanent redirect, request with GET method
			if method != "GET" {
				// Temporary redirect, request with same method
				// As of Go 1.3, Go does not support status code 308.
				code = 307
			}

			if tsr && r.RedirectTrailingSlash {
				var uri string
				if len(path) > 1 && path[len(path)-1] == '/' {
					uri = path[:len(path)-1]
				} else {
					uri = path + "/"
				}

				if len(ctx.URI().QueryString()) > 0 {
					uri += "?" + string(ctx.QueryArgs().QueryString())
				}

				ctx.Redirect(uri, code)
				return
			}

			// Try to fix the request path
			if r.RedirectFixedPath {
				fixedPath, found := root.findCaseInsensitivePath(
					CleanPath(path),
					r.RedirectTrailingSlash,
				)

				if found {
					queryBuf := ctx.URI().QueryString()
					if len(queryBuf) > 0 {
						fixedPath = append(fixedPath, questionMark...)
						fixedPath = append(fixedPath, queryBuf...)
					}
					uri := string(fixedPath)
					ctx.Redirect(uri, code)
					return
				}
			}
		}
	}

	if method == "OPTIONS" {
		// Handle OPTIONS requests
		if r.HandleOPTIONS {
			if allow := r.allowed(path, method); len(allow) > 0 {
				ctx.Response.Header.Set("Allow", allow)
				return
			}
		}
	} else {
		// Handle 405
		if r.HandleMethodNotAllowed {
			if allow := r.allowed(path, method); len(allow) > 0 {
				ctx.Response.Header.Set("Allow", allow)
				if r.MethodNotAllowed != nil {
					r.MethodNotAllowed(ctx)
				} else {
					ctx.SetStatusCode(fasthttp.StatusMethodNotAllowed)
					ctx.SetContentTypeBytes(defaultContentType)
					ctx.SetBodyString(fasthttp.StatusMessage(fasthttp.StatusMethodNotAllowed))
				}
				return
			}
		}
	}

	// Handle 404
	if r.NotFound != nil {
		r.NotFound(ctx)
	} else {
		ctx.Error(fasthttp.StatusMessage(fasthttp.StatusNotFound),
			fasthttp.StatusNotFound)
	}
}
<commit_msg>Using original path in router<commit_after>// Copyright 2013 Julien Schmidt. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be found
// in the LICENSE file.

// Package fasthttprouter is a trie based high performance HTTP request router.
//
// A trivial example is:
//
// package main

// import (
//     "fmt"
//     "log"
//
//     "github.com/buaazp/fasthttprouter"
//     "github.com/valyala/fasthttp"
// )

// func Index(ctx *fasthttp.RequestCtx) {
//     fmt.Fprint(ctx, "Welcome!\n")
// }

// func Hello(ctx *fasthttp.RequestCtx) {
//     fmt.Fprintf(ctx, "hello, %s!\n", ctx.UserValue("name"))
// }

// func main() {
//     router := fasthttprouter.New()
//     router.GET("/", Index)
//     router.GET("/hello/:name", Hello)

//     log.Fatal(fasthttp.ListenAndServe(":8080", router.Handler))
// }
//
// The router matches incoming requests by the request method and the path.
// If a handle is registered for this path and method, the router delegates the
// request to that function.
// For the methods GET, POST, PUT, PATCH and DELETE shortcut functions exist to
// register handles, for all other methods router.Handle can be used.
//
// The registered path, against which the router matches incoming requests, can
// contain two types of parameters:
//  Syntax    Type
//  :name     named parameter
//  *name     catch-all parameter
//
// Named parameters are dynamic path segments. They match anything until the
// next '/' or the path end:
//  Path: /blog/:category/:post
//
//  Requests:
//   /blog/go/request-routers            match: category="go", post="request-routers"
//   /blog/go/request-routers/           no match, but the router would redirect
//   /blog/go/                           no match
//   /blog/go/request-routers/comments   no match
//
// Catch-all parameters match anything until the path end, including the
// directory index (the '/' before the catch-all). Since they match anything
// until the end, catch-all parameters must always be the final path element.
//  Path: /files/*filepath
//
//  Requests:
//   /files/                             match: filepath="/"
//   /files/LICENSE                      match: filepath="/LICENSE"
//   /files/templates/article.html       match: filepath="/templates/article.html"
//   /files                              no match, but the router would redirect
//
// The value of parameters is inside ctx.UserValue
// To retrieve the value of a parameter:
//  // use the name of the parameter
//  user := ps.UserValue("user")
//

package fasthttprouter

import (
	"strings"

	"github.com/valyala/fasthttp"
)

var (
	defaultContentType = []byte("text/plain; charset=utf-8")
	questionMark       = []byte("?")
)

// Router is a http.Handler which can be used to dispatch requests to different
// handler functions via configurable routes
type Router struct {
	trees map[string]*node

	// Enables automatic redirection if the current route can't be matched but a
	// handler for the path with (without) the trailing slash exists.
	// For example if /foo/ is requested but a route only exists for /foo, the
	// client is redirected to /foo with http status code 301 for GET requests
	// and 307 for all other request methods.
	RedirectTrailingSlash bool

	// If enabled, the router tries to fix the current request path, if no
	// handle is registered for it.
	// First superfluous path elements like ../ or // are removed.
	// Afterwards the router does a case-insensitive lookup of the cleaned path.
	// If a handle can be found for this route, the router makes a redirection
	// to the corrected path with status code 301 for GET requests and 307 for
	// all other request methods.
	// For example /FOO and /..//Foo could be redirected to /foo.
	// RedirectTrailingSlash is independent of this option.
	RedirectFixedPath bool

	// If enabled, the router checks if another method is allowed for the
	// current route, if the current request can not be routed.
	// If this is the case, the request is answered with 'Method Not Allowed'
	// and HTTP status code 405.
	// If no other Method is allowed, the request is delegated to the NotFound
	// handler.
	HandleMethodNotAllowed bool

	// If enabled, the router automatically replies to OPTIONS requests.
	// Custom OPTIONS handlers take priority over automatic replies.
	HandleOPTIONS bool

	// Configurable http.Handler which is called when no matching route is
	// found. If it is not set, http.NotFound is used.
	NotFound fasthttp.RequestHandler

	// Configurable http.Handler which is called when a request
	// cannot be routed and HandleMethodNotAllowed is true.
	// If it is not set, http.Error with http.StatusMethodNotAllowed is used.
	// The "Allow" header with allowed request methods is set before the handler
	// is called.
	MethodNotAllowed fasthttp.RequestHandler

	// Function to handle panics recovered from http handlers.
	// It should be used to generate a error page and return the http error code
	// 500 (Internal Server Error).
	// The handler can be used to keep your server from crashing because of
	// unrecovered panics.
	PanicHandler func(*fasthttp.RequestCtx, interface{})
}

// New returns a new initialized Router.
// Path auto-correction, including trailing slashes, is enabled by default.
func New() *Router {
	return &Router{
		RedirectTrailingSlash:  true,
		RedirectFixedPath:      true,
		HandleMethodNotAllowed: true,
		HandleOPTIONS:          true,
	}
}

// GET is a shortcut for router.Handle("GET", path, handle)
func (r *Router) GET(path string, handle fasthttp.RequestHandler) {
	r.Handle("GET", path, handle)
}

// HEAD is a shortcut for router.Handle("HEAD", path, handle)
func (r *Router) HEAD(path string, handle fasthttp.RequestHandler) {
	r.Handle("HEAD", path, handle)
}

// OPTIONS is a shortcut for router.Handle("OPTIONS", path, handle)
func (r *Router) OPTIONS(path string, handle fasthttp.RequestHandler) {
	r.Handle("OPTIONS", path, handle)
}

// POST is a shortcut for router.Handle("POST", path, handle)
func (r *Router) POST(path string, handle fasthttp.RequestHandler) {
	r.Handle("POST", path, handle)
}

// PUT is a shortcut for router.Handle("PUT", path, handle)
func (r *Router) PUT(path string, handle fasthttp.RequestHandler) {
	r.Handle("PUT", path, handle)
}

// PATCH is a shortcut for router.Handle("PATCH", path, handle)
func (r *Router) PATCH(path string, handle fasthttp.RequestHandler) {
	r.Handle("PATCH", path, handle)
}

// DELETE is a shortcut for router.Handle("DELETE", path, handle)
func (r *Router) DELETE(path string, handle fasthttp.RequestHandler) {
	r.Handle("DELETE", path, handle)
}

// Handle registers a new request handle with the given path and method.
//
// For GET, POST, PUT, PATCH and DELETE requests the respective shortcut
// functions can be used.
//
// This function is intended for bulk loading and to allow the usage of less
// frequently used, non-standardized or custom methods (e.g. for internal
// communication with a proxy).
func (r *Router) Handle(method, path string, handle fasthttp.RequestHandler) {
	if path[0] != '/' {
		panic("path must begin with '/' in path '" + path + "'")
	}

	if r.trees == nil {
		r.trees = make(map[string]*node)
	}

	root := r.trees[method]
	if root == nil {
		root = new(node)
		r.trees[method] = root
	}

	root.addRoute(path, handle)
}

// ServeFiles serves files from the given file system root.
// The path must end with "/*filepath", files are then served from the local
// path /defined/root/dir/*filepath.
// For example if root is "/etc" and *filepath is "passwd", the local file
// "/etc/passwd" would be served.
// Internally a http.FileServer is used, therefore http.NotFound is used instead
// of the Router's NotFound handler.
//     router.ServeFiles("/src/*filepath", "/var/www")
func (r *Router) ServeFiles(path string, rootPath string) {
	if len(path) < 10 || path[len(path)-10:] != "/*filepath" {
		panic("path must end with /*filepath in path '" + path + "'")
	}
	prefix := path[:len(path)-10]

	fileHandler := fasthttp.FSHandler(rootPath, strings.Count(prefix, "/"))

	r.GET(path, func(ctx *fasthttp.RequestCtx) {
		fileHandler(ctx)
	})
}

func (r *Router) recv(ctx *fasthttp.RequestCtx) {
	if rcv := recover(); rcv != nil {
		r.PanicHandler(ctx, rcv)
	}
}

// Lookup allows the manual lookup of a method + path combo.
// This is e.g. useful to build a framework around this router.
// If the path was found, it returns the handle function and the path parameter
// values. Otherwise the third return value indicates whether a redirection to
// the same path with an extra / without the trailing slash should be performed.
func (r *Router) Lookup(method, path string, ctx *fasthttp.RequestCtx) (fasthttp.RequestHandler, bool) {
	if root := r.trees[method]; root != nil {
		return root.getValue(path, ctx)
	}
	return nil, false
}

func (r *Router) allowed(path, reqMethod string) (allow string) {
	if path == "*" || path == "/*" { // server-wide
		for method := range r.trees {
			if method == "OPTIONS" {
				continue
			}

			// add request method to list of allowed methods
			if len(allow) == 0 {
				allow = method
			} else {
				allow += ", " + method
			}
		}
	} else { // specific path
		for method := range r.trees {
			// Skip the requested method - we already tried this one
			if method == reqMethod || method == "OPTIONS" {
				continue
			}

			handle, _ := r.trees[method].getValue(path, nil)
			if handle != nil {
				// add request method to list of allowed methods
				if len(allow) == 0 {
					allow = method
				} else {
					allow += ", " + method
				}
			}
		}
	}
	if len(allow) > 0 {
		allow += ", OPTIONS"
	}
	return
}

// Handler makes the router implement the fasthttp.ListenAndServe interface.
func (r *Router) Handler(ctx *fasthttp.RequestCtx) {
	if r.PanicHandler != nil {
		defer r.recv(ctx)
	}

	path := string(ctx.URI().PathOriginal())
	method := string(ctx.Method())
	if root := r.trees[method]; root != nil {
		if f, tsr := root.getValue(path, ctx); f != nil {
			f(ctx)
			return
		} else if method != "CONNECT" && path != "/" {
			code := 301 // Permanent redirect, request with GET method
			if method != "GET" {
				// Temporary redirect, request with same method
				// As of Go 1.3, Go does not support status code 308.
				code = 307
			}

			if tsr && r.RedirectTrailingSlash {
				var uri string
				if len(path) > 1 && path[len(path)-1] == '/' {
					uri = path[:len(path)-1]
				} else {
					uri = path + "/"
				}

				if len(ctx.URI().QueryString()) > 0 {
					uri += "?" + string(ctx.QueryArgs().QueryString())
				}

				ctx.Redirect(uri, code)
				return
			}

			// Try to fix the request path
			if r.RedirectFixedPath {
				fixedPath, found := root.findCaseInsensitivePath(
					CleanPath(path),
					r.RedirectTrailingSlash,
				)

				if found {
					queryBuf := ctx.URI().QueryString()
					if len(queryBuf) > 0 {
						fixedPath = append(fixedPath, questionMark...)
						fixedPath = append(fixedPath, queryBuf...)
					}
					uri := string(fixedPath)
					ctx.Redirect(uri, code)
					return
				}
			}
		}
	}

	if method == "OPTIONS" {
		// Handle OPTIONS requests
		if r.HandleOPTIONS {
			if allow := r.allowed(path, method); len(allow) > 0 {
				ctx.Response.Header.Set("Allow", allow)
				return
			}
		}
	} else {
		// Handle 405
		if r.HandleMethodNotAllowed {
			if allow := r.allowed(path, method); len(allow) > 0 {
				ctx.Response.Header.Set("Allow", allow)
				if r.MethodNotAllowed != nil {
					r.MethodNotAllowed(ctx)
				} else {
					ctx.SetStatusCode(fasthttp.StatusMethodNotAllowed)
					ctx.SetContentTypeBytes(defaultContentType)
					ctx.SetBodyString(fasthttp.StatusMessage(fasthttp.StatusMethodNotAllowed))
				}
				return
			}
		}
	}

	// Handle 404
	if r.NotFound != nil {
		r.NotFound(ctx)
	} else {
		ctx.Error(fasthttp.StatusMessage(fasthttp.StatusNotFound),
			fasthttp.StatusNotFound)
	}
}
<|endoftext|>


<commit_before>package kocha

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"net/http"
	"os"
	"path/filepath"
	"reflect"
	"regexp"
	"strconv"
	"strings"
)

type (
	MethodArgs map[string]string
)

type RouteTable []*Route

type Route struct {
	Name        string
	Path        string
	Controller  interface{}
	MethodTypes map[string]MethodArgs
	RegexpPath  *regexp.Regexp
}

var (
	controllerMethods = map[string]bool{
		"Get":    true,
		"Post":   true,
		"Put":    true,
		"Delete": true,
		"Head":   true,
		"Patch":  true,
	}
	typeRegexpStrings = map[string]string{
		"":    `[\w-]+`, // default
		"int": `\d+`,
	}
	placeHolderRegexp = regexp.MustCompile(`:[\w-]+`)
	pathRegexp        = regexp.MustCompile(`/(?:(?::([\w-]+))|[\w-]*)`)
)

func InitRouteTable(routeTable RouteTable) RouteTable {
	for _, route := range routeTable {
		route.buildMethodTypes()
		route.buildRegexpPath()
	}
	return routeTable
}

func Reverse(name string, v ...interface{}) string {
	for _, route := range appConfig.RouteTable {
		if route.Name == name {
			return route.reverse(v...)
		}
	}
	types := make([]string, len(v))
	for i, value := range v {
		types[i] = reflect.TypeOf(value).Name()
	}
	panic(fmt.Errorf("no match route found: %v (%v)", name, strings.Join(types, ", ")))
}

func dispatch(req *http.Request) (controller *reflect.Value, method *reflect.Value, args []reflect.Value) {
	methodName := strings.Title(strings.ToLower(req.Method))
	path := normPath(req.URL.Path)
	for _, route := range appConfig.RouteTable {
		if controller, method, args = route.dispatch(methodName, path); controller != nil {
			break
		}
	}
	return controller, method, args
}

func (route *Route) dispatch(methodName, path string) (controller *reflect.Value, method *reflect.Value, args []reflect.Value) {
	matchesBase := route.RegexpPath.FindStringSubmatch(path)
	if len(matchesBase) == 0 {
		return nil, nil, nil
	}
	matchesBase = matchesBase[1:]
	matches := make([]string, 0, len(matchesBase))
	subexpNames := route.RegexpPath.SubexpNames()[1:]
	types := make([]string, 0, len(subexpNames))
	if args, ok := route.MethodTypes[methodName]; ok {
		for name, t := range args {
			for i, subexpName := range subexpNames {
				if subexpName == name {
					matches = append(matches, matchesBase[i])
					types = append(types, t)
					break
				}
			}
		}
	}
	for i, v := range matches {
		switch types[i] {
		case "int":
			p, err := strconv.Atoi(v)
			if err != nil {
				panic(err)
			}
			args = append(args, reflect.ValueOf(p))
		default:
			args = append(args, reflect.ValueOf(v))
		}
	}
	t := reflect.TypeOf(route.Controller)
	c := reflect.New(t)
	m := c.MethodByName(methodName)
	return &c, &m, args
}

func (route *Route) reverse(v ...interface{}) string {
	switch n := route.RegexpPath.NumSubexp(); {
	case len(v) < n:
		panic(fmt.Errorf("too few arguments: %v (controller is %v)", route.Name, reflect.TypeOf(route.Controller).Name()))
	case len(v) > n:
		panic(fmt.Errorf("too many arguments: %v (controller is %v)", route.Name, reflect.TypeOf(route.Controller).Name()))
	case len(v)+n == 0:
		return route.Path
	}
	path := placeHolderRegexp.ReplaceAllStringFunc(route.Path, func(s string) string {
		result := fmt.Sprint(v[0])
		v = v[1:]
		return result
	})
	if !route.RegexpPath.MatchString(path) {
		panic(fmt.Errorf("parameter type mismatch: %v (controller is %v)", route.Name, reflect.TypeOf(route.Controller).Name()))
	}
	return path
}

func (route *Route) buildMethodTypes() {
	controller := reflect.TypeOf(route.Controller)
	cname := controller.Name()
	sname := ToSnakeCase(cname)
	pkgPath := controller.PkgPath()
	goFile := sname + ".go"
	var filePath string
	for _, src := range build.Default.SrcDirs() {
		if path, err := filepath.Abs(filepath.Join(src, pkgPath, goFile)); err == nil {
			if _, err := os.Stat(path); err == nil {
				filePath = path
				break
			}
		}
	}
	if filePath == "" {
		panic(fmt.Errorf("%s: no such file", filepath.Join(pkgPath, goFile)))
	}
	f, err := parser.ParseFile(token.NewFileSet(), filePath, nil, 0)
	if err != nil {
		panic(err)
	}
	route.MethodTypes = make(map[string]MethodArgs)
	for _, d := range f.Decls {
		ast.Inspect(d, func(node ast.Node) bool {
			fdecl, ok := node.(*ast.FuncDecl)
			if !ok {
				return false
			}
			var recv string
			switch t := fdecl.Recv.List[0].Type.(type) {
			case *ast.StarExpr:
				recv = t.X.(*ast.Ident).Name
			case *ast.Ident:
				recv = t.Name
			}
			if recv != cname {
				return false
			}
			methodName := fdecl.Name.Name
			if _, ok := controllerMethods[methodName]; !ok {
				return false
			}
			route.MethodTypes[methodName] = make(MethodArgs)
			for _, v := range fdecl.Type.Params.List {
				t := v.Type.(*ast.Ident).Name
				for _, name := range v.Names {
					route.MethodTypes[methodName][name.Name] = t
				}
			}
			return false
		})
	}
}

func (route *Route) buildRegexpPath() {
	var regexpBuf bytes.Buffer
	for _, paths := range pathRegexp.FindAllStringSubmatch(route.Path, -1) {
		name := paths[1]
		if name == "" {
			regexpBuf.WriteString(regexp.QuoteMeta(paths[0]))
			continue
		}
		var rePatStr string
		for _, args := range route.MethodTypes {
			if t, ok := args[name]; ok {
				if rePatStr = typeRegexpStrings[t]; rePatStr == "" {
					rePatStr = typeRegexpStrings[""]
				}
				break
			}
		}
		if rePatStr == "" {
			methodNames := make([]string, 0, len(route.MethodTypes))
			for methodName, _ := range route.MethodTypes {
				methodNames = append(methodNames, methodName)
			}
			controllerName := reflect.TypeOf(route.Controller).Name()
			panic(fmt.Errorf("argument `%s` is not defined in these methods `%s.%s`",
				name, controllerName, strings.Join(methodNames, ", ")))
		}
		regexpBuf.WriteString(fmt.Sprintf(`/(?P<%s>%s)`, regexp.QuoteMeta(name), rePatStr))
	}
	route.RegexpPath = regexp.MustCompile(fmt.Sprintf("^%s$", regexpBuf.String()))
}
<commit_msg>Improve the build to method types.<commit_after>package kocha

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"net/http"
	"os"
	"path/filepath"
	"reflect"
	"regexp"
	"strconv"
	"strings"
)

type (
	MethodArgs map[string]string
)

type RouteTable []*Route

type Route struct {
	Name        string
	Path        string
	Controller  interface{}
	MethodTypes map[string]MethodArgs
	RegexpPath  *regexp.Regexp
}

var (
	controllerMethods = map[string]bool{
		"Get":    true,
		"Post":   true,
		"Put":    true,
		"Delete": true,
		"Head":   true,
		"Patch":  true,
	}
	typeRegexpStrings = map[string]string{
		"":    `[\w-]+`, // default
		"int": `\d+`,
	}
	placeHolderRegexp = regexp.MustCompile(`:[\w-]+`)
	pathRegexp        = regexp.MustCompile(`/(?:(?::([\w-]+))|[\w-]*)`)
)

func InitRouteTable(routeTable RouteTable) RouteTable {
	for _, route := range routeTable {
		route.buildMethodTypes()
		route.buildRegexpPath()
	}
	return routeTable
}

func Reverse(name string, v ...interface{}) string {
	for _, route := range appConfig.RouteTable {
		if route.Name == name {
			return route.reverse(v...)
		}
	}
	types := make([]string, len(v))
	for i, value := range v {
		types[i] = reflect.TypeOf(value).Name()
	}
	panic(fmt.Errorf("no match route found: %v (%v)", name, strings.Join(types, ", ")))
}

func dispatch(req *http.Request) (controller *reflect.Value, method *reflect.Value, args []reflect.Value) {
	methodName := strings.Title(strings.ToLower(req.Method))
	path := normPath(req.URL.Path)
	for _, route := range appConfig.RouteTable {
		if controller, method, args = route.dispatch(methodName, path); controller != nil {
			break
		}
	}
	return controller, method, args
}

func (route *Route) dispatch(methodName, path string) (controller *reflect.Value, method *reflect.Value, args []reflect.Value) {
	matchesBase := route.RegexpPath.FindStringSubmatch(path)
	if len(matchesBase) == 0 {
		return nil, nil, nil
	}
	matchesBase = matchesBase[1:]
	matches := make([]string, 0, len(matchesBase))
	subexpNames := route.RegexpPath.SubexpNames()[1:]
	types := make([]string, 0, len(subexpNames))
	if args, ok := route.MethodTypes[methodName]; ok {
		for name, t := range args {
			for i, subexpName := range subexpNames {
				if subexpName == name {
					matches = append(matches, matchesBase[i])
					types = append(types, t)
					break
				}
			}
		}
	}
	for i, v := range matches {
		switch types[i] {
		case "int":
			p, err := strconv.Atoi(v)
			if err != nil {
				panic(err)
			}
			args = append(args, reflect.ValueOf(p))
		default:
			args = append(args, reflect.ValueOf(v))
		}
	}
	t := reflect.TypeOf(route.Controller)
	c := reflect.New(t)
	m := c.MethodByName(methodName)
	return &c, &m, args
}

func (route *Route) reverse(v ...interface{}) string {
	switch n := route.RegexpPath.NumSubexp(); {
	case len(v) < n:
		panic(fmt.Errorf("too few arguments: %v (controller is %v)", route.Name, reflect.TypeOf(route.Controller).Name()))
	case len(v) > n:
		panic(fmt.Errorf("too many arguments: %v (controller is %v)", route.Name, reflect.TypeOf(route.Controller).Name()))
	case len(v)+n == 0:
		return route.Path
	}
	path := placeHolderRegexp.ReplaceAllStringFunc(route.Path, func(s string) string {
		result := fmt.Sprint(v[0])
		v = v[1:]
		return result
	})
	if !route.RegexpPath.MatchString(path) {
		panic(fmt.Errorf("parameter type mismatch: %v (controller is %v)", route.Name, reflect.TypeOf(route.Controller).Name()))
	}
	return path
}

func (route *Route) buildMethodTypes() {
	controller := reflect.TypeOf(route.Controller)
	cname := controller.Name()
	pkgPath := controller.PkgPath()
	pkgDir := findPkgDir(pkgPath)
	if pkgDir == "" {
		panic(fmt.Errorf("%v: package not found", pkgPath))
	}
	pkgInfo, err := build.ImportDir(pkgDir, 0)
	if err != nil {
		panic(err)
	}
	astFiles := make([]*ast.File, len(pkgInfo.GoFiles))
	for i, goFilePath := range pkgInfo.GoFiles {
		if astFiles[i], err = parser.ParseFile(token.NewFileSet(), filepath.Join(pkgInfo.Dir, goFilePath), nil, 0); err != nil {
			panic(err)
		}
	}
	route.MethodTypes = make(map[string]MethodArgs)
	for _, file := range astFiles {
		for _, d := range file.Decls {
			ast.Inspect(d, func(node ast.Node) bool {
				fdecl, ok := node.(*ast.FuncDecl)
				if !ok || fdecl.Recv == nil {
					return false
				}
				var recv string
				switch t := fdecl.Recv.List[0].Type.(type) {
				case *ast.StarExpr:
					recv = t.X.(*ast.Ident).Name
				case *ast.Ident:
					recv = t.Name
				}
				if recv != cname {
					return false
				}
				methodName := fdecl.Name.Name
				if _, ok := controllerMethods[methodName]; !ok {
					return false
				}
				route.MethodTypes[methodName] = make(MethodArgs)
				for _, v := range fdecl.Type.Params.List {
					t := v.Type.(*ast.Ident).Name
					for _, name := range v.Names {
						route.MethodTypes[methodName][name.Name] = t
					}
				}
				return false
			})
		}
	}
}

func findPkgDir(pkgPath string) string {
	var pkgDir string
	for _, srcDir := range build.Default.SrcDirs() {
		path, err := filepath.Abs(filepath.Join(srcDir, pkgPath))
		if err != nil {
			panic(err)
		}
		if _, err := os.Stat(path); err == nil {
			pkgDir = path
			break
		}
	}
	return pkgDir
}

func (route *Route) buildRegexpPath() {
	var regexpBuf bytes.Buffer
	for _, paths := range pathRegexp.FindAllStringSubmatch(route.Path, -1) {
		name := paths[1]
		if name == "" {
			regexpBuf.WriteString(regexp.QuoteMeta(paths[0]))
			continue
		}
		var rePatStr string
		for _, args := range route.MethodTypes {
			if t, ok := args[name]; ok {
				if rePatStr = typeRegexpStrings[t]; rePatStr == "" {
					rePatStr = typeRegexpStrings[""]
				}
				break
			}
		}
		if rePatStr == "" {
			methodNames := make([]string, 0, len(route.MethodTypes))
			for methodName, _ := range route.MethodTypes {
				methodNames = append(methodNames, methodName)
			}
			controllerName := reflect.TypeOf(route.Controller).Name()
			panic(fmt.Errorf("argument `%s` is not defined in these methods `%s.%s`",
				name, controllerName, strings.Join(methodNames, ", ")))
		}
		regexpBuf.WriteString(fmt.Sprintf(`/(?P<%s>%s)`, regexp.QuoteMeta(name), rePatStr))
	}
	route.RegexpPath = regexp.MustCompile(fmt.Sprintf("^%s$", regexpBuf.String()))
}
<|endoftext|>


<commit_before>package router

import (
	"github.com/karlseguin/params"
	"github.com/karlseguin/scratch"
	"net/http"
	"strings"
)

type Handler func(out http.ResponseWriter, req *Request)

type Router struct {
	notFound  Handler
	routes    map[string]*RoutePart
	paramPool *params.Pool
	valuePool *scratch.StringsPool
}

func New(config *Configuration) *Router {
	router := &Router{
		notFound: notFoundHandler,
		routes: map[string]*RoutePart{
			"GET":     newRoutePart(),
			"POST":    newRoutePart(),
			"PUT":     newRoutePart(),
			"DELETE":  newRoutePart(),
			"PURGE":   newRoutePart(),
			"PATCH":   newRoutePart(),
			"HEAD":    newRoutePart(),
			"OPTIONS": newRoutePart(),
		},
	}
	router.paramPool = params.NewPool(config.paramPoolSize, config.paramPoolCount)
	router.valuePool = scratch.NewStrings(config.paramPoolSize, config.paramPoolCount)
	return router
}

func (r *Router) NotFound(handler Handler) {
	r.notFound = handler
}

func (r *Router) All(path string, handler Handler) {
	for _, rp := range r.routes {
		r.add(rp, path, handler)
	}
}

func (r *Router) Get(path string, handler Handler) {
	r.add(r.routes["GET"], path, handler)
}

func (r *Router) Post(path string, handler Handler) {
	r.add(r.routes["POST"], path, handler)
}

func (r *Router) Put(path string, handler Handler) {
	r.add(r.routes["PUT"], path, handler)
}

func (r *Router) Delete(path string, handler Handler) {
	r.add(r.routes["DELETE"], path, handler)
}

func (r *Router) Purge(path string, handler Handler) {
	r.add(r.routes["PURGE"], path, handler)
}

func (r *Router) Patch(path string, handler Handler) {
	r.add(r.routes["PATCH"], path, handler)
}

func (r *Router) Options(path string, handler Handler) {
	r.add(r.routes["OPTIONS"], path, handler)
}

func (r *Router) ServeHTTP(out http.ResponseWriter, httpReq *http.Request) {
	req := &Request{Request: httpReq, params: params.Empty}
	rp, ok := r.routes[req.Method]
	if ok == false {
		r.notFound(out, req)
		return
	}
	path := req.URL.Path
	if path == "" || path == "/" {
		handler := rp.handler
		if handler == nil {
			handler = r.notFound
		}
		handler(out, req)
		return
	}
	if path[0] == '/' {
		path = path[1:]
	}

	if path[len(path)-1] == '/' {
		path = path[:(len(path) - 1)]
	}

	values := r.valuePool.Checkout()
	defer values.Release()

	var handler Handler
	for {
		original := rp
		index := strings.Index(path, "/")
		if index == -1 {
			index = len(path)
		}
		part := path[:index]
		if rp, ok = rp.parts[part]; ok == false {
			if original.prefixes != nil {
				lower := strings.ToLower(part)
				for _, prefix := range original.prefixes {
					if len(prefix.value) == 0 || strings.HasPrefix(lower, prefix.value) {
						rp = original
						handler = prefix.handler
						break
					}
				}
				if handler != nil {
					break
				}
			}
			if rp, ok = original.parts[":"]; ok == false {
				break
			}
			values.Add(part)
		}

		if rp == nil || len(path) == index {
			break
		}
		path = path[index+1:]
	}

	if rp == nil || (rp.handler == nil && handler == nil) {
		r.notFound(out, req)
		return
	}
	if l := values.Len(); l > 0 {
		params := r.paramPool.Checkout()
		defer params.Release()
		for i, value := range values.Values() {
			params.Set(rp.params[i], value)
		}
		req.params = params
	}
	if handler == nil {
		handler = rp.handler
	}
	handler(out, req)

}

func (r *Router) add(rp *RoutePart, path string, handler Handler) {
	if path == "" || path == "/" {
		rp.handler = handler
		return
	}

	if path[0] == '/' {
		path = path[1:]
	}
	if path[len(path)-1] == '/' {
		path = path[:(len(path) - 1)]
	}

	params := make([]string, 0, 1)
	parts := strings.Split(path, "/")
	for _, part := range parts {
		if part[len(part)-1] == '*' {
			if rp.prefixes == nil {
				rp.prefixes = make([]*Prefix, 0, 1)
			}
			prefix := &Prefix{value: strings.ToLower(part[:len(part)-1]), handler: handler}
			rp.prefixes = appendPrefix(rp.prefixes, prefix)
			break
		}
		if part[0] == ':' {
			params = appendOne(params, part[1:])
			part = ":"
		}
		sub, exists := rp.parts[part]
		if exists == false {
			sub = newRoutePart()
			rp.parts[part] = sub
		}
		rp = sub
	}
	if len(params) > 0 {
		rp.params = params
	}
	if rp.handler == nil {
		rp.handler = handler
	}
}

func (r Router) Routes() map[string]*RoutePart {
	return r.routes
}

func notFoundHandler(out http.ResponseWriter, req *Request) {
	out.WriteHeader(404)
}

func appendOne(arr []string, value string) []string {
	target := arr
	if len(arr) == cap(arr) {
		target = make([]string, len(arr)+1)
		copy(target, arr)
	}
	return append(target, value)
}

func appendPrefix(arr []*Prefix, value *Prefix) []*Prefix {
	target := arr
	if len(arr) == cap(arr) {
		target = make([]*Prefix, len(arr)+1)
		copy(target, arr)
	}
	return append(target, value)
}
<commit_msg>Expose router.Lookup()  which can be used to lookup a route and extract its parameters without going through the ServeHTTP process. This allows router to be used outside of the http workflow.<commit_after>package router

import (
	"github.com/karlseguin/params"
	"github.com/karlseguin/scratch"
	"net/http"
	"strings"
)

type Handler func(out http.ResponseWriter, req *Request)

type Router struct {
	notFound  Handler
	routes    map[string]*RoutePart
	paramPool *params.Pool
	valuePool *scratch.StringsPool
}

func New(config *Configuration) *Router {
	router := &Router{
		notFound: notFoundHandler,
		routes: map[string]*RoutePart{
			"GET":     newRoutePart(),
			"POST":    newRoutePart(),
			"PUT":     newRoutePart(),
			"DELETE":  newRoutePart(),
			"PURGE":   newRoutePart(),
			"PATCH":   newRoutePart(),
			"HEAD":    newRoutePart(),
			"OPTIONS": newRoutePart(),
		},
	}
	router.paramPool = params.NewPool(config.paramPoolSize, config.paramPoolCount)
	router.valuePool = scratch.NewStrings(config.paramPoolSize, config.paramPoolCount)
	return router
}

func (r *Router) NotFound(handler Handler) {
	r.notFound = handler
}

func (r *Router) All(path string, handler Handler) {
	for _, rp := range r.routes {
		r.add(rp, path, handler)
	}
}

func (r *Router) Get(path string, handler Handler) {
	r.add(r.routes["GET"], path, handler)
}

func (r *Router) Post(path string, handler Handler) {
	r.add(r.routes["POST"], path, handler)
}

func (r *Router) Put(path string, handler Handler) {
	r.add(r.routes["PUT"], path, handler)
}

func (r *Router) Delete(path string, handler Handler) {
	r.add(r.routes["DELETE"], path, handler)
}

func (r *Router) Purge(path string, handler Handler) {
	r.add(r.routes["PURGE"], path, handler)
}

func (r *Router) Patch(path string, handler Handler) {
	r.add(r.routes["PATCH"], path, handler)
}

func (r *Router) Options(path string, handler Handler) {
	r.add(r.routes["OPTIONS"], path, handler)
}

func (r *Router) ServeHTTP(out http.ResponseWriter, hr *http.Request) {
	params, handler, _ := r.Lookup(hr)
	defer params.Release()
	req := &Request{Request: hr, params: params}
	if handler == nil {
		handler = r.notFound
	}
	handler(out, req)
}

func (r *Router) Lookup(req *http.Request) (params.Params, Handler, *RoutePart) {
	rp, ok := r.routes[req.Method]
	var params params.Params = params.Empty
	if ok == false {
		return params, nil, nil
	}
	path := req.URL.Path
	if path == "" || path == "/" {
		handler := rp.handler
		if handler == nil {
			handler = r.notFound
		}
		return params, handler, rp
	}

	if path[0] == '/' {
		path = path[1:]
	}
	if path[len(path)-1] == '/' {
		path = path[:(len(path) - 1)]
	}

	values := r.valuePool.Checkout()
	defer values.Release()

	var handler Handler
	for {
		original := rp
		index := strings.Index(path, "/")
		if index == -1 {
			index = len(path)
		}
		part := path[:index]
		if rp, ok = rp.parts[part]; ok == false {
			if original.prefixes != nil {
				lower := strings.ToLower(part)
				for _, prefix := range original.prefixes {
					if len(prefix.value) == 0 || strings.HasPrefix(lower, prefix.value) {
						rp = original
						handler = prefix.handler
						break
					}
				}
				if handler != nil {
					break
				}
			}
			if rp, ok = original.parts[":"]; ok == false {
				break
			}
			values.Add(part)
		}

		if rp == nil || len(path) == index {
			break
		}
		path = path[index+1:]
	}

	if rp == nil || (rp.handler == nil && handler == nil) {
		return params, nil, nil
	}

	if l := values.Len(); l > 0 {
		params = r.paramPool.Checkout()
		for i, value := range values.Values() {
			params.Set(rp.params[i], value)
		}
	}
	if handler == nil {
		handler = rp.handler
	}
	return params, handler, rp
}

func (r *Router) add(rp *RoutePart, path string, handler Handler) {
	if path == "" || path == "/" {
		rp.handler = handler
		return
	}

	if path[0] == '/' {
		path = path[1:]
	}
	if path[len(path)-1] == '/' {
		path = path[:(len(path) - 1)]
	}

	params := make([]string, 0, 1)
	parts := strings.Split(path, "/")
	for _, part := range parts {
		if part[len(part)-1] == '*' {
			if rp.prefixes == nil {
				rp.prefixes = make([]*Prefix, 0, 1)
			}
			prefix := &Prefix{value: strings.ToLower(part[:len(part)-1]), handler: handler}
			rp.prefixes = appendPrefix(rp.prefixes, prefix)
			break
		}
		if part[0] == ':' {
			params = appendOne(params, part[1:])
			part = ":"
		}
		sub, exists := rp.parts[part]
		if exists == false {
			sub = newRoutePart()
			rp.parts[part] = sub
		}
		rp = sub
	}
	if len(params) > 0 {
		rp.params = params
	}
	if rp.handler == nil {
		rp.handler = handler
	}
}

func (r Router) Routes() map[string]*RoutePart {
	return r.routes
}

func notFoundHandler(out http.ResponseWriter, req *Request) {
	out.WriteHeader(404)
}

func appendOne(arr []string, value string) []string {
	target := arr
	if len(arr) == cap(arr) {
		target = make([]string, len(arr)+1)
		copy(target, arr)
	}
	return append(target, value)
}

func appendPrefix(arr []*Prefix, value *Prefix) []*Prefix {
	target := arr
	if len(arr) == cap(arr) {
		target = make([]*Prefix, len(arr)+1)
		copy(target, arr)
	}
	return append(target, value)
}
<|endoftext|>


<commit_before>package easyroute

import (
	"github.com/gorilla/mux"
	"net/http"
	"time"
)

type handlerFunc func(*Request)
type beforeHandlerFunc func(*Request) bool
type loggerFunc func(string, ...interface{})

type Logger struct {
	LogI loggerFunc
	LogE loggerFunc
	LogD loggerFunc
}

type Router struct {
	// Inherit a mux router
	*mux.Router

	beforeHandler beforeHandlerFunc
	logger        Logger
}

// NewRouter creates a new easyroute Router object with the provided
// before handler and logger struct
func NewRouter(beforeFn beforeHandlerFunc, logger Logger) Router {
	muxRouter := mux.NewRouter()

	router := Router{
		muxRouter,
		beforeFn,
		logger,
	}

	return router
}

// SubRoute creates a new easyroute Router off the base router with provided
// prefix. This preserves the same before handler.
func (g *Router) SubRoute(prefix string) Router {
	muxRouter := g.PathPrefix(prefix).Subrouter()

	router := Router{
		muxRouter,
		g.beforeHandler,
		g.logger,
	}

	return router
}

// SubRouteC creates a new easyroute Router off the base router with provided
// prefix and an additional before handler.
// The routes in this router will now run through first the parent(base) router's
// before handler and then this router's before handler.
func (g *Router) SubRouteC(prefix string, beforeFn beforeHandlerFunc) Router {
	muxRouter := g.PathPrefix(prefix).Subrouter()

	router := Router{
		muxRouter,
		func(r *Request) bool {
			if g.beforeHandler(r) {
				return beforeFn(r)
			}
			return false
		},
		g.logger,
	}

	return router
}

func (g *Router) requestHandler(fn handlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		var body interface{}
		// Start timer
		start := time.Now()

		request := NewRequest(w, r)
		// Run Before block
		if g.beforeHandler(&request) == true {
			// If the before block returns false we don't execute the rest
			// Run actual handler
			fn(&request)
		}

		if g.logger.LogI != nil {
			// Log out some of the info
			origin := r.RemoteAddr
			method := r.Method
			path := r.URL.Path
			userUuid := request.UserUuid
			request.Body(body)
			elapsed := time.Since(start)
			g.logger.LogI("origin=%s method=%s path=%s body=%s user_uuid=%s elapsed=%s", origin, method, path, body, userUuid, elapsed)
		}
	}
}

func (g *Router) Get(path string, handler handlerFunc) {
	g.HandleFunc(path, g.requestHandler(handler)).Methods("GET")
}

func (g *Router) Put(path string, handler handlerFunc) {
	g.HandleFunc(path, g.requestHandler(handler)).Methods("PUT")
}

func (g *Router) Post(path string, handler handlerFunc) {
	g.HandleFunc(path, g.requestHandler(handler)).Methods("POST")
}

func (g *Router) Delete(path string, handler handlerFunc) {
	g.HandleFunc(path, g.requestHandler(handler)).Methods("DELETE")
}
<commit_msg>integrated airbrake into middleware<commit_after>package easyroute

import (
	"github.com/gorilla/mux"
	"net/http"
	"time"
	"strconv"
	gobrake "gopkg.in/airbrake/gobrake.v2"
)

type handlerFunc func(*Request)
type beforeHandlerFunc func(*Request) bool
type loggerFunc func(string, ...interface{})

type Logger struct {
	LogI loggerFunc
	LogE loggerFunc
	LogD loggerFunc
}

type Router struct {
	// Inherit a mux router
	*mux.Router

	beforeHandler beforeHandlerFunc
	logger        Logger

	airbrakeProjectId int64
	airbrakeProjectKey string
	airbrakeEnabled bool
}

// NewRouter creates a new easyroute Router object with the provided
// before handler and logger struct
func NewRouter(beforeFn beforeHandlerFunc, logger Logger, airbrakeInfo map[string]string) Router {
	muxRouter := mux.NewRouter()
	var enableAirbrake bool
	if(airbrakeInfo["enable"] == "true"){
		enableAirbrake = true
	} else {
		enableAirbrake = false
	}
	airbrakeId, err := strconv.Atoi(airbrakeInfo["projectId"])
	if(err != nil){
		enableAirbrake = false
	}
	router := Router{
		muxRouter,
		beforeFn,
		logger,
		int64(airbrakeId),
		airbrakeInfo["projectKey"],
		enableAirbrake,
	}

	return router
}

// SubRoute creates a new easyroute Router off the base router with provided
// prefix. This preserves the same before handler.
func (g *Router) SubRoute(prefix string) Router {
	muxRouter := g.PathPrefix(prefix).Subrouter()

	router := Router{
		muxRouter,
		g.beforeHandler,
		g.logger,
		0,
		"",
		false,
	}

	return router
}

// SubRouteC creates a new easyroute Router off the base router with provided
// prefix and an additional before handler.
// The routes in this router will now run through first the parent(base) router's
// before handler and then this router's before handler.
func (g *Router) SubRouteC(prefix string, beforeFn beforeHandlerFunc) Router {
	muxRouter := g.PathPrefix(prefix).Subrouter()

	router := Router{
		muxRouter,
		func(r *Request) bool {
			if g.beforeHandler(r) {
				return beforeFn(r)
			}
			return false
		},
		g.logger,
		0,
		"",
		false,
	}

	return router
}

func (g *Router) requestHandler(fn handlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if(g.airbrakeEnabled == true){
			airbrake := gobrake.NewNotifier(g.airbrakeProjectId, g.airbrakeProjectKey)
			defer airbrake.Close()
			defer airbrake.NotifyOnPanic()
		}
		var body interface{}
		// Start timer
		start := time.Now()

		request := NewRequest(w, r)
		// Run Before block
		if g.beforeHandler(&request) == true {
			// If the before block returns false we don't execute the rest
			// Run actual handler
			fn(&request)
		}

		if g.logger.LogI != nil {
			// Log out some of the info
			origin := r.RemoteAddr
			method := r.Method
			path := r.URL.Path
			userUuid := request.UserUuid
			request.Body(body)
			elapsed := time.Since(start)
			g.logger.LogI("origin=%s method=%s path=%s body=%s user_uuid=%s elapsed=%s", origin, method, path, body, userUuid, elapsed)
		}
	}
}

func (g *Router) Get(path string, handler handlerFunc) {
	g.HandleFunc(path, g.requestHandler(handler)).Methods("GET")
}

func (g *Router) Put(path string, handler handlerFunc) {
	g.HandleFunc(path, g.requestHandler(handler)).Methods("PUT")
}

func (g *Router) Post(path string, handler handlerFunc) {
	g.HandleFunc(path, g.requestHandler(handler)).Methods("POST")
}

func (g *Router) Delete(path string, handler handlerFunc) {
	g.HandleFunc(path, g.requestHandler(handler)).Methods("DELETE")
}
<|endoftext|>


<commit_before>package node

import (
	"sync"
)

// A router handles all routing tasks that don't involve the local machine
// including connection management, reachability handling, packet receipt
// relaying, and (outstanding) payment tracking. AKA anything which doesn't
// need the private key.
type router struct {
	pk PublicKey
	// A map of public key hashes to connections
	connections map[NodeAddress]Connection

	*reachabilityHandler
	*routingHandler
	*receiptHandler
	*ledger

	lock *sync.Mutex
	quit chan bool
}

func newRouter(pk PublicKey) *router {
	reach := newReachability(pk.Hash())
	routing := newRouting(pk, reach)
	c1, c2, quit := splitChannel(routing.Routes())
	receipt := newReceipt(pk.Hash(), c1)
	ledger := newLedger(pk.Hash(), receipt.PacketHashes(), c2)
	return &router{
		pk,
		make(map[NodeAddress]Connection),
		reach,
		routing,
		receipt,
		ledger,
		&sync.Mutex{},
		quit,
	}
}

func (r *router) GetAddress() PublicKey {
	return r.pk
}

func (r *router) Connections() []Connection {
	r.lock.Lock()
	defer r.lock.Unlock()
	c := make([]Connection, 0, len(r.connections))
	for _, k := range r.connections {
		c = append(c, k)
	}
	return c
}

func (r *router) AddConnection(c Connection) {
	r.lock.Lock()
	defer r.lock.Unlock()
	id := c.Key().Hash()
	_, duplicate := r.connections[id]
	if duplicate {
		c.Close()
		return
	}
	r.connections[id] = c

	// Curry the id since the various sub connections don't know about it
	r.routingHandler.AddConnection(id, c)
	r.reachabilityHandler.AddConnection(id, c)
	r.receiptHandler.AddConnection(id, c)
}

func (r *router) Close() error {
	r.reachabilityHandler.Close()
	r.routingHandler.Close()
	r.receiptHandler.Close()
	r.ledger.Close()
	close(r.quit)
	return nil
}
<commit_msg>Add connection counter expvar<commit_after>package node

import (
	"expvar"
	"fmt"
	"sync"
)

var connections_export *expvar.Map

func init() {
	if connections_export == nil {
		connections_export = expvar.NewMap("connections")
	}
}

// A router handles all routing tasks that don't involve the local machine
// including connection management, reachability handling, packet receipt
// relaying, and (outstanding) payment tracking. AKA anything which doesn't
// need the private key.
type router struct {
	pk PublicKey
	// A map of public key hashes to connections
	connections map[NodeAddress]Connection

	*reachabilityHandler
	*routingHandler
	*receiptHandler
	*ledger

	lock *sync.Mutex
	quit chan bool
}

func newRouter(pk PublicKey) *router {
	reach := newReachability(pk.Hash())
	routing := newRouting(pk, reach)
	c1, c2, quit := splitChannel(routing.Routes())
	receipt := newReceipt(pk.Hash(), c1)
	ledger := newLedger(pk.Hash(), receipt.PacketHashes(), c2)
	return &router{
		pk,
		make(map[NodeAddress]Connection),
		reach,
		routing,
		receipt,
		ledger,
		&sync.Mutex{},
		quit,
	}
}

func (r *router) GetAddress() PublicKey {
	return r.pk
}

func (r *router) Connections() []Connection {
	r.lock.Lock()
	defer r.lock.Unlock()
	c := make([]Connection, 0, len(r.connections))
	for _, k := range r.connections {
		c = append(c, k)
	}
	return c
}

func (r *router) AddConnection(c Connection) {
	r.lock.Lock()
	defer r.lock.Unlock()
	id := c.Key().Hash()
	_, duplicate := r.connections[id]
	if duplicate {
		c.Close()
		return
	}
	r.connections[id] = c

	// Curry the id since the various sub connections don't know about it
	r.routingHandler.AddConnection(id, c)
	r.reachabilityHandler.AddConnection(id, c)
	r.receiptHandler.AddConnection(id, c)
	connections_export.Add(fmt.Sprintf("%x", c.Key().Hash()), 1)
}

func (r *router) Close() error {
	r.reachabilityHandler.Close()
	r.routingHandler.Close()
	r.receiptHandler.Close()
	r.ledger.Close()
	close(r.quit)
	return nil
}
<|endoftext|>


<commit_before>package dtls

import (
	"net"
	"reflect"
	"sync"
	"testing"
	"time"

	"github.com/pion/dtls/v2/internal/net/dpipe"
	"github.com/pion/transport/test"
)

func TestReplayProtection(t *testing.T) {
	// Limit runtime in case of deadlocks
	lim := test.TimeOut(5 * time.Second)
	defer lim.Stop()

	// Check for leaking routines
	report := test.CheckRoutines(t)
	defer report()

	c0, c1 := dpipe.Pipe()
	c2, c3 := dpipe.Pipe()
	conn := []net.Conn{c0, c1, c2, c3}

	var wgReplays, wgRoutines sync.WaitGroup

	replayer := func(ca, cb net.Conn) {
		defer wgRoutines.Done()
		// Man in the middle
		for {
			b := make([]byte, 2048)
			n, rerr := ca.Read(b)
			if rerr != nil {
				return
			}
			if _, werr := cb.Write(b[:n]); werr != nil {
				t.Error(werr)
				return
			}

			wgReplays.Add(1)
			go func() {
				defer wgReplays.Done()
				// Replay bit later
				time.Sleep(time.Millisecond)
				if _, werr := cb.Write(b[:n]); werr != nil {
					t.Error(werr)
				}
			}()
		}
	}
	wgRoutines.Add(2)
	go replayer(conn[1], conn[2])
	go replayer(conn[2], conn[1])

	ca, cb, err := pipeConn(conn[0], conn[3])
	if err != nil {
		t.Fatal(err)
	}

	const numMsgs = 10

	var received [2][][]byte
	for i, c := range []net.Conn{ca, cb} {
		i := i
		c := c
		wgRoutines.Add(1)
		wgReplays.Add(1) // Keep locked until the final message
		var lastMsgDone sync.Once
		go func() {
			defer wgRoutines.Done()
			for {
				b := make([]byte, 2048)
				n, rerr := c.Read(b)
				if rerr != nil {
					return
				}
				received[i] = append(received[i], b[:n])
				if b[0] == numMsgs-1 {
					// Final message received
					lastMsgDone.Do(func() {
						wgReplays.Done()
					})
				}
			}
		}()
	}

	var sent [][]byte
	for i := 0; i < numMsgs; i++ {
		data := []byte{byte(i)}
		sent = append(sent, data)
		if _, werr := ca.Write(data); werr != nil {
			t.Error(werr)
			return
		}
		if _, werr := cb.Write(data); werr != nil {
			t.Error(werr)
			return
		}
	}
	wgReplays.Wait()
	time.Sleep(10 * time.Millisecond) // Ensure all replayed packets are sent

	for i := 0; i < 4; i++ {
		if err := conn[i].Close(); err != nil {
			t.Error(err)
		}
	}
	if err := ca.Close(); err != nil {
		t.Error(err)
	}
	if err := cb.Close(); err != nil {
		t.Error(err)
	}
	wgRoutines.Wait()

	for _, r := range received {
		if !reflect.DeepEqual(sent, r) {
			t.Errorf("Received data differs, expected: %v, got: %v", sent, r)
		}
	}
}
<commit_msg>Fix data race in TestReplayProtection<commit_after>package dtls

import (
	"context"
	"net"
	"reflect"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/pion/dtls/v2/internal/net/dpipe"
	"github.com/pion/transport/test"
)

func TestReplayProtection(t *testing.T) {
	// Limit runtime in case of deadlocks
	lim := test.TimeOut(5 * time.Second)
	defer lim.Stop()

	// Check for leaking routines
	report := test.CheckRoutines(t)
	defer report()

	c0, c1 := dpipe.Pipe()
	c2, c3 := dpipe.Pipe()
	conn := []net.Conn{c0, c1, c2, c3}

	var wgRoutines sync.WaitGroup
	var cntReplays int32 = 1

	ctxReplayDone, replayDone := context.WithCancel(context.Background())

	replaySendDone := func() {
		cnt := atomic.AddInt32(&cntReplays, -1)
		if cnt == 0 {
			replayDone()
		}
	}

	replayer := func(ca, cb net.Conn) {
		defer wgRoutines.Done()
		// Man in the middle
		for {
			b := make([]byte, 2048)
			n, rerr := ca.Read(b)
			if rerr != nil {
				return
			}
			if _, werr := cb.Write(b[:n]); werr != nil {
				t.Error(werr)
				return
			}

			atomic.AddInt32(&cntReplays, 1)
			go func() {
				defer replaySendDone()
				// Replay bit later
				time.Sleep(time.Millisecond)
				if _, werr := cb.Write(b[:n]); werr != nil {
					t.Error(werr)
				}
			}()
		}
	}
	wgRoutines.Add(2)
	go replayer(conn[1], conn[2])
	go replayer(conn[2], conn[1])

	ca, cb, err := pipeConn(conn[0], conn[3])
	if err != nil {
		t.Fatal(err)
	}

	const numMsgs = 10

	var received [2][][]byte
	for i, c := range []net.Conn{ca, cb} {
		i := i
		c := c
		wgRoutines.Add(1)
		atomic.AddInt32(&cntReplays, 1) // Keep locked until the final message
		var lastMsgDone sync.Once
		go func() {
			defer wgRoutines.Done()
			for {
				b := make([]byte, 2048)
				n, rerr := c.Read(b)
				if rerr != nil {
					return
				}
				received[i] = append(received[i], b[:n])
				if b[0] == numMsgs-1 {
					// Final message received
					lastMsgDone.Do(func() {
						defer replaySendDone()
					})
				}
			}
		}()
	}

	var sent [][]byte
	for i := 0; i < numMsgs; i++ {
		data := []byte{byte(i)}
		sent = append(sent, data)
		if _, werr := ca.Write(data); werr != nil {
			t.Error(werr)
			return
		}
		if _, werr := cb.Write(data); werr != nil {
			t.Error(werr)
			return
		}
	}

	replaySendDone()
	<-ctxReplayDone.Done()
	time.Sleep(10 * time.Millisecond) // Ensure all replayed packets are sent

	for i := 0; i < 4; i++ {
		if err := conn[i].Close(); err != nil {
			t.Error(err)
		}
	}
	if err := ca.Close(); err != nil {
		t.Error(err)
	}
	if err := cb.Close(); err != nil {
		t.Error(err)
	}
	wgRoutines.Wait()

	for _, r := range received {
		if !reflect.DeepEqual(sent, r) {
			t.Errorf("Received data differs, expected: %v, got: %v", sent, r)
		}
	}
}
<|endoftext|>