readlink_test.go

package os_test

import (
	"encoding/binary"
	"io/ioutil"
	"os"
	"runtime"
	"syscall"
	"unsafe"

	"testing"
)

var (
	modadvapi32 = syscall.NewLazyDLL("advapi32.dll")
	modkernel32 = syscall.NewLazyDLL("kernel32.dll")

	procAdjustTokenPrivileges = modadvapi32.NewProc("AdjustTokenPrivileges")
	procImpersonateSelf       = modadvapi32.NewProc("ImpersonateSelf")
	procLookupPrivilegeValue  = modadvapi32.NewProc("LookupPrivilegeValueW")
	procOpenThreadToken       = modadvapi32.NewProc("OpenThreadToken")
	procRevertToSelf          = modadvapi32.NewProc("RevertToSelf")

	procGetCurrentThread = modkernel32.NewProc("GetCurrentThread")

	le = binary.LittleEndian
)

const (
	SE_CREATE_SYMBOLIC_LINK_NAME = "SeCreateSymbolicLinkPrivilege"
	SE_RESTORE_NAME              = "SeRestorePrivilege"

	SE_PRIVILEGE_ENABLED = 2

	SYMLINK_FLAG_RELATIVE = 1

	FSCTL_SET_REPARSE_POINT = 0x900a4

	securityImpersonation = 2
)

func TestReadlink(t *testing.T) {
	defer chtmpdir(t)()

	f, err := os.Create("test")
	if err != nil {
		t.Fatal(err)
	}
	f.Close()

	err = Symlink("test", "linkToTest")
	if err != nil {
		t.Fatal(err)
	}

	defer func() {
		if r := recover(); r != nil {
			t.Fatal(r)
		}
	}()

	target, err := os.Readlink("linkToTest")
	if err != nil {
		t.Fatal(err)
	}

	if target != "test" {
		t.Error("readlink is broken")
	}
}

func Symlink(oldname string, newname string) (err error) {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	{
		err = impersonateSelf(securityImpersonation)
		if err != nil {
			return err
		}
		defer revertToSelf()

		th, err := getCurrentThread()
		if err != nil {
			return err
		}

		var t syscall.Token

		err = openThreadToken(th, syscall.TOKEN_ADJUST_PRIVILEGES|syscall.TOKEN_QUERY, false, &t)
		if err != nil {
			return err
		}
		defer t.Close()

		err = enablePrivileges(t, []string{SE_CREATE_SYMBOLIC_LINK_NAME, SE_RESTORE_NAME})
		if err != nil {
			return err
		}
	}

	fd, err := syscall.CreateFile(syscall.StringToUTF16Ptr(newname),
		syscall.GENERIC_READ|syscall.GENERIC_WRITE, 0, nil, syscall.CREATE_ALWAYS,
		syscall.FILE_FLAG_OPEN_REPARSE_POINT|syscall.FILE_FLAG_BACKUP_SEMANTICS, 0)
	if err != nil {
		return err
	}
	defer syscall.CloseHandle(fd)

	sympath := syscall.StringToUTF16(oldname)

	sympath = sympath[:len(sympath)-1]

	// typedef struct _REPARSE_DATA_BUFFER {
	//   uint32 ReparseTag;
	//   uint16 ReparseDataLength;
	//   uint16 Reserved;
	//   uint16 SubstituteNameOffset;
	//   uint16 SubstituteNameLength;
	//   uint16 PrintNameOffset;
	//   uint16 PrintNameLength;
	//   uint32  Flags;
	//   WCHAR  PathBuffer[1];
	// } REPARSE_DATA_BUFFER, *PREPARSE_DATA_BUFFER;

	rdbbuf := make([]byte, 20+len(sympath)*4)

	le.PutUint32(rdbbuf[:4], syscall.IO_REPARSE_TAG_SYMLINK)
	le.PutUint16(rdbbuf[4:6], uint16(len(rdbbuf)-8))
	le.PutUint16(rdbbuf[8:10], 0)
	le.PutUint16(rdbbuf[10:12], uint16(len(sympath)*2))
	for i, w := range sympath {
		le.PutUint16(rdbbuf[20+i*2:20+i*2+2], w)
	}
	le.PutUint16(rdbbuf[12:14], uint16(len(sympath)*2))
	le.PutUint16(rdbbuf[14:16], uint16(len(sympath)*2))
	le.PutUint16(rdbbuf[16:20], SYMLINK_FLAG_RELATIVE)
	for i, w := range sympath {
		le.PutUint16(rdbbuf[20+len(sympath)*2+i*2:20+len(sympath)*2+i*2+2], w)
	}

	var returnBytes uint32 // TODO it needs invastigation. without this variable, it panics on windows 7. why? is it known bug?

	return syscall.DeviceIoControl(fd, FSCTL_SET_REPARSE_POINT, &rdbbuf[0], uint32(len(rdbbuf)), nil, 0, &returnBytes, nil)
	// return syscall.DeviceIoControl(fd, FSCTL_SET_REPARSE_POINT, &rdbbuf[0], uint32(len(rdbbuf)), nil, 0, nil, nil)
}

func enablePrivileges(t syscall.Token, names []string) error {
	// type TokenPrivileges struct {
	// 	PrivilegeCount uint32
	// 	Privileges     []struct {
	// 		Luid       uint64
	// 		Attributes uint32
	// 	}
	// }

	tp := make([]byte, 4+12*len(names))

	le.PutUint32(tp[:4], uint32(len(names)))

	off := 4

	var luid uint64

	for _, name := range names {
		err := lookupPrivilegeValue("", name, &luid)
		if err != nil {
			return err
		}
		le.PutUint64(tp[off:off+8], luid)
		le.PutUint32(tp[off+8:off+12], SE_PRIVILEGE_ENABLED)

		off += 12
	}

	err := adjustTokenPrivileges(t, false, &tp[0], 0, nil, nil)
	if err != nil {
		return err
	}

	return nil
}

func lookupPrivilegeValue(systemName string, name string, luid *uint64) (err error) {
	var _p0 *uint16
	_p0, err = syscall.UTF16PtrFromString(systemName)
	if err != nil {
		return
	}
	var _p1 *uint16
	_p1, err = syscall.UTF16PtrFromString(name)
	if err != nil {
		return
	}

	r1, _, e1 := syscall.Syscall(procLookupPrivilegeValue.Addr(), 3, uintptr(unsafe.Pointer(_p0)), uintptr(unsafe.Pointer(_p1)), uintptr(unsafe.Pointer(luid)))
	if r1 == 0 {
		if e1 != 0 {
			err = error(e1)
		} else {
			err = syscall.EINVAL
		}
	}
	return
}

func adjustTokenPrivileges(token syscall.Token, releaseAll bool, input *byte, outputSize uint32, output *byte, requiredSize *uint32) (err error) {
	var _p0 uint32
	if releaseAll {
		_p0 = 1
	} else {
		_p0 = 0
	}
	r0, _, e1 := syscall.Syscall6(procAdjustTokenPrivileges.Addr(), 6, uintptr(token), uintptr(_p0), uintptr(unsafe.Pointer(input)), uintptr(outputSize), uintptr(unsafe.Pointer(output)), uintptr(unsafe.Pointer(requiredSize)))
	if r0 == 0 {
		if e1 != 0 {
			err = error(e1)
		} else {
			err = syscall.EINVAL
		}
	}
	return
}

func openThreadToken(thread syscall.Handle, accessMask uint32, openAsSelf bool, token *syscall.Token) (err error) {
	var _p0 uint32
	if openAsSelf {
		_p0 = 1
	} else {
		_p0 = 0
	}
	r1, _, e1 := syscall.Syscall6(procOpenThreadToken.Addr(), 4, uintptr(thread), uintptr(accessMask), uintptr(_p0), uintptr(unsafe.Pointer(token)), 0, 0)
	if r1 == 0 {
		if e1 != 0 {
			err = error(e1)
		} else {
			err = syscall.EINVAL
		}
	}
	return
}

func getCurrentThread() (pseudoHandle syscall.Handle, err error) {
	r0, _, e1 := syscall.Syscall(procGetCurrentThread.Addr(), 0, 0, 0, 0)
	pseudoHandle = syscall.Handle(r0)
	if pseudoHandle == 0 {
		if e1 != 0 {
			err = error(e1)
		} else {
			err = syscall.EINVAL
		}
	}
	return
}

func impersonateSelf(level uint32) (err error) {
	r1, _, e1 := syscall.Syscall(procImpersonateSelf.Addr(), 1, uintptr(level), 0, 0)
	if r1 == 0 {
		if e1 != 0 {
			err = error(e1)
		} else {
			err = syscall.EINVAL
		}
	}
	return
}

func revertToSelf() (err error) {
	r1, _, e1 := syscall.Syscall(procRevertToSelf.Addr(), 0, 0, 0, 0)
	if r1 == 0 {
		if e1 != 0 {
			err = error(e1)
		} else {
			err = syscall.EINVAL
		}
	}
	return
}

// chtmpdir changes the working directory to a new temporary directory and
// provides a cleanup function. Used when PWD is read-only.
func chtmpdir(t *testing.T) func() {
	oldwd, err := os.Getwd()
	if err != nil {
		t.Fatalf("chtmpdir: %v", err)
	}
	d, err := ioutil.TempDir("", "test")
	if err != nil {
		t.Fatalf("chtmpdir: %v", err)
	}
	if err := os.Chdir(d); err != nil {
		t.Fatalf("chtmpdir: %v", err)
	}
	return func() {
		if err := os.Chdir(oldwd); err != nil {
			t.Fatalf("chtmpdir: %v", err)
		}
		os.RemoveAll(d)
	}
}