<commit_before>// Copyright (C) 2014 <NAME> <<EMAIL>>
// License: MIT

// Package nbd uses the Linux NBD layer to emulate a block device in user space
package nbd

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"runtime"
	"sync"
	"syscall"
)

const (
	// Defined in <linux/fs.h>:
	BLKROSET = 4701

	// Defined in <linux/nbd.h>:
	ioctlSetSock       = 43776
	ioctlSetBlockSize  = 43777
	ioctlDoIt          = 43779
	ioctlClearSock     = 43780
	ioctlClearQueue    = 43781
	ioctlSetSizeBlocks = 43783
	ioctlDisconnect    = 43784
	ioctlSetFlags      = 43786
)

const (
	cmdRead  = 0
	cmdWrite = 1
	cmdDisc  = 2
	cmdFlush = 3
	cmdTrim  = 4
)

const (
	flagSendFlush = (1 << 2) // can flush writeback cache
	flagSendTrim  = (1 << 5) // Send TRIM (discard)
	// flagHasFlags   = (1 << 0) // nbd-server supports flags
	// flagReadOnly   = (1 << 1) // device is read-only
	// flagSendFUA    = (1 << 3) // Send FUA (Force Unit Access)
	// flagRotational = (1 << 4) // Use elevator algorithm - rotational media
)

const (
	magicRequest = 0x25609513
	magicReply   = 0x67446698
	// Do *not* use magics: 0x12560953 0x96744668.
)

const (
	errIO = 5
)

// ioctl() helper function
func ioctl(a1, a2, a3 uintptr) (err error) {
	_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, a1, a2, a3)
	if errno != 0 {
		err = errno
	}
	return err
}

// Device interface is a subset of os.File.
type Device interface {
	ReadAt(b []byte, off int64) (n int, err error)
	WriteAt(b []byte, off int64) (n int, err error)
	Sync() error
	Trim(off, len int64) error
}

type NBD struct {
	device    Device
	size      int64
	blocksize int64
	nbd       *os.File
	socket    int
	setsocket int
	closer    chan error
}

func Create(device Device, size int64, blocksize int64) *NBD {
	if size >= 0 {
		return &NBD{
			device:    device,
			size:      size,
			blocksize: blocksize,
			nbd:       nil,
			socket:    0,
		}
	}
	return nil
}

// return true if connected
func (nbd *NBD) IsConnected() bool {
	return nbd.nbd != nil && nbd.socket > 0
}

func (nbd *NBD) Size() int64 {
	return nbd.size
}

func (nbd *NBD) SetSize(size, blocksize int64) error {
	if err := ioctl(nbd.nbd.Fd(), ioctlSetBlockSize, uintptr(blocksize)); err != nil {
		return &os.PathError{
			Path: nbd.nbd.Name(),
			Op:   "ioctl NBD_SET_BLKSIZE",
			Err:  err,
		}
	}
	if err := ioctl(nbd.nbd.Fd(), ioctlSetSizeBlocks, uintptr(size/blocksize)); err != nil {
		return &os.PathError{
			Path: nbd.nbd.Name(),
			Op:   "ioctl NBD_SET_SIZE_BLOCKS",
			Err:  err,
		}
	}
	return nil
}

func FindDevice() (string, error) {
	// FIXME: Oh god... fixme.
	// find free nbd device
	for i := 0; ; i++ {
		dev := fmt.Sprintf("/dev/nbd%d", i)
		if _, err := os.Stat(dev); os.IsNotExist(err) {
			break // no more devices
		}
		if _, err := os.Stat(fmt.Sprintf("/sys/block/nbd%d/pid", i)); !os.IsNotExist(err) {
			continue // busy
		}
		return dev, nil
	}
	return "", errors.New("no devices available")
}

func (nbd *NBD) OpenDevice(dev string) (string, error) {
	f, err := os.Open(dev)
	if err != nil {
		return "", err
	}
	nbd.nbd = f

	// possible candidate
	ioctl(f.Fd(), BLKROSET, 0) // I'm really sorry about this
	pair, err := syscall.Socketpair(syscall.SOCK_STREAM, syscall.AF_UNIX, 0)
	if err != nil {
		return "", err
	}
	if err := ioctl(f.Fd(), ioctlSetSock, uintptr(pair[0])); err != nil {
		return "", err
	}

	nbd.setsocket = pair[0] // FIXME: We shouldn't hold on to this.
	nbd.socket = pair[1]
	return dev, nil
}

func (nbd *NBD) Serve() error {
	if err := nbd.SetSize(nbd.size, nbd.blocksize); err != nil {
		return err // already set by nbd.Size()
	}
	if err := ioctl(nbd.nbd.Fd(), ioctlSetFlags, uintptr(flagSendFlush|flagSendTrim)); err != nil {
		return &os.PathError{
			Path: nbd.nbd.Name(),
			Op:   "ioctl NBD_SET_FLAGS",
			Err:  err,
		}
	}

	c := &serverConn{
		rw: os.NewFile(uintptr(nbd.socket), "<nbd socket>"),
	}
	n := runtime.GOMAXPROCS(0)

	wg := new(sync.WaitGroup)
	wg.Add(n)
	for i := 0; i < n; i++ {
		go c.serveLoop(nbd.device, wg)
	}
	if err := waitDisconnect(nbd.nbd); err != nil {
		return err
	}
	wg.Wait()
	return nil<commit_msg>add TODO for handling multithreading<commit_after>// Copyright (C) 2014 <NAME> <<EMAIL>>
// License: MIT

// Package nbd uses the Linux NBD layer to emulate a block device in user space
package nbd

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"runtime"
	"sync"
	"syscall"
)

const (
	// Defined in <linux/fs.h>:
	BLKROSET = 4701

	// Defined in <linux/nbd.h>:
	ioctlSetSock       = 43776
	ioctlSetBlockSize  = 43777
	ioctlDoIt          = 43779
	ioctlClearSock     = 43780
	ioctlClearQueue    = 43781
	ioctlSetSizeBlocks = 43783
	ioctlDisconnect    = 43784
	ioctlSetFlags      = 43786
)

const (
	cmdRead  = 0
	cmdWrite = 1
	cmdDisc  = 2
	cmdFlush = 3
	cmdTrim  = 4
)

const (
	flagSendFlush = (1 << 2) // can flush writeback cache
	flagSendTrim  = (1 << 5) // Send TRIM (discard)
	// flagHasFlags   = (1 << 0) // nbd-server supports flags
	// flagReadOnly   = (1 << 1) // device is read-only
	// flagSendFUA    = (1 << 3) // Send FUA (Force Unit Access)
	// flagRotational = (1 << 4) // Use elevator algorithm - rotational media
)

const (
	magicRequest = 0x25609513
	magicReply   = 0x67446698
	// Do *not* use magics: 0x12560953 0x96744668.
)

const (
	errIO = 5
)

// ioctl() helper function
func ioctl(a1, a2, a3 uintptr) (err error) {
	_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, a1, a2, a3)
	if errno != 0 {
		err = errno
	}
	return err
}

// Device interface is a subset of os.File.
type Device interface {
	ReadAt(b []byte, off int64) (n int, err error)
	WriteAt(b []byte, off int64) (n int, err error)
	Sync() error
	Trim(off, len int64) error
}

type NBD struct {
	device    Device
	size      int64
	blocksize int64
	nbd       *os.File
	socket    int
	setsocket int
	closer    chan error
}

func Create(device Device, size int64, blocksize int64) *NBD {
	if size >= 0 {
		return &NBD{
			device:    device,
			size:      size,
			blocksize: blocksize,
			nbd:       nil,
			socket:    0,
		}
	}
	return nil
}

// return true if connected
func (nbd *NBD) IsConnected() bool {
	return nbd.nbd != nil && nbd.socket > 0
}

func (nbd *NBD) Size() int64 {
	return nbd.size
}

func (nbd *NBD) SetSize(size, blocksize int64) error {
	if err := ioctl(nbd.nbd.Fd(), ioctlSetBlockSize, uintptr(blocksize)); err != nil {
		return &os.PathError{
			Path: nbd.nbd.Name(),
			Op:   "ioctl NBD_SET_BLKSIZE",
			Err:  err,
		}
	}
	if err := ioctl(nbd.nbd.Fd(), ioctlSetSizeBlocks, uintptr(size/blocksize)); err != nil {
		return &os.PathError{
			Path: nbd.nbd.Name(),
			Op:   "ioctl NBD_SET_SIZE_BLOCKS",
			Err:  err,
		}
	}
	return nil
}

func FindDevice() (string, error) {
	// FIXME: Oh god... fixme.
	// find free nbd device
	for i := 0; ; i++ {
		dev := fmt.Sprintf("/dev/nbd%d", i)
		if _, err := os.Stat(dev); os.IsNotExist(err) {
			break // no more devices
		}
		if _, err := os.Stat(fmt.Sprintf("/sys/block/nbd%d/pid", i)); !os.IsNotExist(err) {
			continue // busy
		}
		return dev, nil
	}
	return "", errors.New("no devices available")
}

func (nbd *NBD) OpenDevice(dev string) (string, error) {
	f, err := os.Open(dev)
	if err != nil {
		return "", err
	}
	nbd.nbd = f

	// possible candidate
	ioctl(f.Fd(), BLKROSET, 0) // I'm really sorry about this
	pair, err := syscall.Socketpair(syscall.SOCK_STREAM, syscall.AF_UNIX, 0)
	if err != nil {
		return "", err
	}
	if err := ioctl(f.Fd(), ioctlSetSock, uintptr(pair[0])); err != nil {
		return "", err
	}

	nbd.setsocket = pair[0] // FIXME: We shouldn't hold on to this.
	nbd.socket = pair[1]
	return dev, nil
}

func (nbd *NBD) Serve() error {
	if err := nbd.SetSize(nbd.size, nbd.blocksize); err != nil {
		return err // already set by nbd.Size()
	}
	if err := ioctl(nbd.nbd.Fd(), ioctlSetFlags, uintptr(flagSendFlush|flagSendTrim)); err != nil {
		return &os.PathError{
			Path: nbd.nbd.Name(),
			Op:   "ioctl NBD_SET_FLAGS",
			Err:  err,
		}
	}

	c := &serverConn{
		rw: os.NewFile(uintptr(nbd.socket), "<nbd socket>"),
	}
	// TODO(barakmich): Scale up NBD by handling multiple requests.
	// Requires thread-safety across the block.BlockFile/agro.File
	//n := runtime.GOMAXPROCS(0) - 1
	n := 1

	wg := new(sync.WaitGroup)
	wg.Add(n)
	for i := 0; i < n; i++ {
		go c.serveLoop(nbd.device, wg)
	}
	if err := waitDisconnect(nbd.nbd); err != nil {
		return err
	}
	wg.Wait()
	return nil
{'fraction_non_alphanumeric': 0.1486034853186918, 'fraction_numerical': 0.03318214370971592, 'mean_word_length': 2.9598298676748582, 'pattern_counts': {'":': 0, '<': 41, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6094870', 'n_tokens_mistral': 3705, 'n_tokens_neox': 3351, 'n_words': 919}