WorkstationService.cpp [plain text]
#include <dns_sd.h> // DNSServiceDiscovery
#include <unistd.h> // usleep
#include <CoreFoundation/CoreFoundation.h>
#include <IOKit/IOKitLib.h>
#include <IOKit/network/IOEthernetInterface.h>
#include <IOKit/network/IONetworkInterface.h>
#include <IOKit/network/IOEthernetController.h>
#include <SystemConfiguration/SystemConfiguration.h>
#include "CLog.h"
#include "WorkstationService.h"
#define kWorkstationType "_workstation._tcp."
#define kWorkstationPort 9
#define kMACAddressLengh (sizeof(" [00:00:00:00:00:00]") - 1)
#define kMaxComputerName 63-kMACAddressLengh
typedef struct MyDNSServiceState {
DNSServiceRef service;
CFSocketRef socket;
} MyDNSServiceState;
DNSServiceErrorType RegisterWorkstationService(MyDNSServiceState *ref, CFStringRef serviceName);
CFStringRef CopyNextWorkstationName(SCDynamicStoreRef store, CFStringRef currentName);
CFStringRef gCurrentWorkstationName = NULL;
MyDNSServiceState *gServiceRegistration = NULL;
CFRunLoopSourceRef gNameMonitorSource = NULL;
static MyDNSServiceState *
MyDNSServiceAlloc(void)
{
MyDNSServiceState *ref = (MyDNSServiceState *) malloc(sizeof(MyDNSServiceState));
assert(ref);
ref->service = NULL;
ref->socket = NULL;
return ref;
}
static void
MyDNSServiceFree(MyDNSServiceState *ref)
{
assert(ref);
if (ref->socket) {
CFSocketInvalidate(ref->socket);
CFRelease(ref->socket);
usleep(1000);
}
if (ref->service) {
DNSServiceRefDeallocate(ref->service);
}
free(ref);
}
static void
MyDNSServiceSocketCallBack(CFSocketRef s, CFSocketCallBackType type, CFDataRef address, const void *data, void *info)
{
#pragma unused(s, type, address, data)
DNSServiceErrorType error;
MyDNSServiceState *ref = (MyDNSServiceState *)info; assert(ref);
error = DNSServiceProcessResult(ref->service);
if (kDNSServiceErr_NoError != error) {
fprintf(stderr, "DNSServiceProcessResult returned %d\n", error);
MyDNSServiceFree(ref);
}
}
static void
MyDNSServiceAddToRunLoop(MyDNSServiceState *ref)
{
CFSocketNativeHandle sock;
CFOptionFlags sockFlags;
CFRunLoopSourceRef source;
CFSocketContext context = { 0, ref, NULL, NULL, NULL };
assert(ref);
sock = DNSServiceRefSockFD(ref->service);
assert(sock >= 0);
ref->socket = CFSocketCreateWithNative(NULL, sock, kCFSocketReadCallBack, MyDNSServiceSocketCallBack, &context);
assert(ref->socket);
sockFlags = CFSocketGetSocketFlags(ref->socket);
CFSocketSetSocketFlags(ref->socket, sockFlags & (~kCFSocketCloseOnInvalidate));
source = CFSocketCreateRunLoopSource(NULL, ref->socket, 0);
assert(source);
CFRunLoopAddSource(CFRunLoopGetCurrent(), source, kCFRunLoopCommonModes);
CFRelease(source);
}
static void
CFStringTruncateToUTF8Length(CFMutableStringRef str, ssize_t utf8LengthLimit)
{
CFIndex shortLen;
CFIndex convertedLen;
CFIndex originalLen;
CFIndex utf8Length;
CFCharacterSetRef whiteCharSet;
UniChar thisChar;
CFIndex trailingCharsToDelete;
originalLen = CFStringGetLength(str);
shortLen = originalLen;
do {
convertedLen = CFStringGetBytes(str, CFRangeMake(0, shortLen), kCFStringEncodingUTF8, 0, false, NULL, 0, &utf8Length);
assert( (convertedLen == shortLen) || (convertedLen == (shortLen - 1)) );
shortLen = convertedLen;
if (utf8Length <= utf8LengthLimit) {
break;
}
shortLen -= 1;
} while (true);
whiteCharSet = CFCharacterSetGetPredefined(kCFCharacterSetWhitespaceAndNewline);
assert(whiteCharSet != NULL);
do {
if ( shortLen == 0 ) {
break;
}
thisChar = CFStringGetCharacterAtIndex(str, shortLen - 1);
if ( ! CFCharacterSetIsCharacterMember(whiteCharSet, thisChar) ) {
break;
}
shortLen -= 1;
} while (true);
trailingCharsToDelete = originalLen - shortLen;
CFStringDelete(str, CFRangeMake(originalLen - trailingCharsToDelete, trailingCharsToDelete));
}
static kern_return_t
FindEthernetInterfaces(io_iterator_t *matchingServices)
{
kern_return_t kernResult;
CFMutableDictionaryRef matchingDict;
CFMutableDictionaryRef propertyMatchDict;
matchingDict = IOServiceMatching(kIOEthernetInterfaceClass);
if (NULL == matchingDict) {
printf("IOServiceMatching returned a NULL dictionary.\n");
} else {
propertyMatchDict = CFDictionaryCreateMutable(kCFAllocatorDefault, 0,
&kCFTypeDictionaryKeyCallBacks,
&kCFTypeDictionaryValueCallBacks);
if (NULL == propertyMatchDict) {
printf("CFDictionaryCreateMutable returned a NULL dictionary.\n");
} else {
CFDictionarySetValue(propertyMatchDict, CFSTR(kIOPrimaryInterface), kCFBooleanTrue);
CFDictionarySetValue(matchingDict, CFSTR(kIOPropertyMatchKey), propertyMatchDict);
CFRelease(propertyMatchDict);
}
}
kernResult = IOServiceGetMatchingServices(kIOMasterPortDefault, matchingDict, matchingServices);
if (KERN_SUCCESS != kernResult) {
printf("IOServiceGetMatchingServices returned 0x%08x\n", kernResult);
}
return kernResult;
}
static kern_return_t
GetMACAddress(io_iterator_t intfIterator, UInt8 *MACAddress, UInt8 bufferSize)
{
io_object_t intfService;
io_object_t controllerService;
kern_return_t kernResult = KERN_FAILURE;
if (bufferSize < kIOEthernetAddressSize) {
return kernResult;
}
bzero(MACAddress, bufferSize);
while ((intfService = IOIteratorNext(intfIterator))) {
CFTypeRef MACAddressAsCFData;
kernResult = IORegistryEntryGetParentEntry(intfService, kIOServicePlane, &controllerService);
if (KERN_SUCCESS != kernResult) {
printf("IORegistryEntryGetParentEntry returned 0x%08x\n", kernResult);
} else {
MACAddressAsCFData = IORegistryEntryCreateCFProperty(controllerService, CFSTR(kIOMACAddress), kCFAllocatorDefault, 0);
if (MACAddressAsCFData) {
CFDataGetBytes((CFDataRef)MACAddressAsCFData, CFRangeMake(0, kIOEthernetAddressSize), MACAddress);
CFRelease(MACAddressAsCFData);
}
(void)IOObjectRelease(controllerService);
}
(void)IOObjectRelease(intfService);
}
return kernResult;
}
static CFStringRef
CopyPrimaryMacAddress(void)
{
io_iterator_t intfIterator;
UInt8 MAC[kIOEthernetAddressSize];
CFStringRef macAddress = NULL;
kern_return_t kernResult = FindEthernetInterfaces(&intfIterator);
if (KERN_SUCCESS != kernResult) {
printf("FindEthernetInterfaces returned 0x%08x\n", kernResult);
} else {
kernResult = GetMACAddress(intfIterator, MAC, sizeof(MAC));
if (KERN_SUCCESS != kernResult) {
printf("GetMACAddress returned 0x%08x\n", kernResult);
} else {
macAddress = CFStringCreateWithFormat(NULL, NULL, CFSTR("[%02x:%02x:%02x:%02x:%02x:%02x]"), MAC[0], MAC[1], MAC[2], MAC[3], MAC[4], MAC[5]);
}
}
(void)IOObjectRelease(intfIterator);
return macAddress;
}
static void
RegisterWorkstationCallBack(DNSServiceRef service, DNSServiceFlags flags, DNSServiceErrorType errorCode,
const char *name, const char *type, const char *domain, void *context)
{
#pragma unused(flags, service, context)
if (errorCode == kDNSServiceErr_NoError) {
printf("Registered Workstation service %s.%s%s\n", name, type, domain);
} else if (errorCode == kDNSServiceErr_NameConflict) {
fprintf(stderr, "ERROR: Workstation name conflict occured so something weird is happening\n");
if (gServiceRegistration) {
MyDNSServiceFree(gServiceRegistration);
gServiceRegistration = NULL;
}
CFStringRef newWorkstationName = CopyNextWorkstationName(NULL, gCurrentWorkstationName);
if (gCurrentWorkstationName) CFRelease(gCurrentWorkstationName);
gCurrentWorkstationName = newWorkstationName;
if (gCurrentWorkstationName) {
gServiceRegistration = MyDNSServiceAlloc();
RegisterWorkstationService(gServiceRegistration, gCurrentWorkstationName);
}
} else {
fprintf(stderr, "ERROR: Workstation name registration failed with error (%d)\n", errorCode);
if (gServiceRegistration) {
MyDNSServiceFree(gServiceRegistration);
gServiceRegistration = NULL;
}
}
}
DNSServiceErrorType
RegisterWorkstationService(MyDNSServiceState *ref, CFStringRef serviceName)
{
char name[64];
DNSServiceErrorType error = kDNSServiceErr_BadParam;
if (CFStringGetCString(serviceName, name, sizeof(name), kCFStringEncodingUTF8)) {
error = DNSServiceRegister(&ref->service,
kDNSServiceFlagsNoAutoRename,
kDNSServiceInterfaceIndexAny,
name,
kWorkstationType,
NULL,
NULL,
htons(kWorkstationPort),
0,
NULL,
RegisterWorkstationCallBack,
(void *)ref);
}
if (kDNSServiceErr_NoError == error) {
MyDNSServiceAddToRunLoop(ref);
}
return error;
}
CFStringRef
CopyNextWorkstationName(SCDynamicStoreRef store, CFStringRef currentName)
{
CFMutableStringRef computerName = NULL;
CFStringRef macAddress = NULL;
if (currentName) {
CFIndex totalLengh = CFStringGetLength(currentName);
if (totalLengh >= (CFIndex)kMACAddressLengh) {
CFRange range = CFRangeMake(0, totalLengh - kMACAddressLengh);
CFStringRef oldComputerName = CFStringCreateWithSubstring(NULL, currentName, range);
if (CFStringHasSuffix(oldComputerName, CFSTR(")"))) {
CFRange result;
if (CFStringFindWithOptions(oldComputerName, CFSTR("("), range, kCFCompareBackwards, &result)) {
range = CFRangeMake(result.location + 1, CFStringGetLength(oldComputerName) - result.location - 2);
CFStringRef countString = CFStringCreateWithSubstring(NULL, currentName, range);
SInt32 conflictCount = CFStringGetIntValue(countString);
if (conflictCount) {
range = CFRangeMake(0, result.location);
CFStringRef tempComputerName = CFStringCreateWithSubstring(NULL, oldComputerName, range);
computerName = CFStringCreateMutableCopy(NULL, 0, tempComputerName);
CFStringRef numberString = CFStringCreateWithFormat(NULL, NULL, CFSTR(" (%d)"), ++conflictCount);
CFStringTruncateToUTF8Length(computerName, kMaxComputerName - CFStringGetLength(numberString));
CFStringAppend(computerName, numberString);
}
}
}
if (!computerName) {
computerName = CFStringCreateMutableCopy(NULL, 0, oldComputerName);
CFStringTruncateToUTF8Length(computerName, kMaxComputerName - 4);
CFStringAppend(computerName, CFSTR(" (2)"));
}
CFRelease(oldComputerName);
} else {
fprintf(stderr, "ERROR: Workstation name is shorter than a MAC address which shouldn't be possible\n");
}
} else {
CFStringRef tempName = SCDynamicStoreCopyComputerName(store, NULL);
if (tempName) {
computerName = CFStringCreateMutableCopy(NULL, 0, tempName);
CFRelease(tempName);
CFStringTruncateToUTF8Length(computerName, kMaxComputerName);
} else {
return NULL;
}
}
macAddress = CopyPrimaryMacAddress();
if (!macAddress) {
if (computerName) {
CFRelease(computerName);
}
return NULL;
}
CFStringAppend(computerName, CFSTR(" "));
CFStringAppend(computerName, macAddress);
CFRelease(macAddress);
return computerName;
}
static void
ComputerNameChangedCallBack(SCDynamicStoreRef store, CFArrayRef changedKeys, void *info)
{
#pragma unused(changedKeys, info)
if (gServiceRegistration) {
MyDNSServiceFree(gServiceRegistration);
gServiceRegistration = NULL;
}
if (gCurrentWorkstationName) {
CFRelease(gCurrentWorkstationName);
}
gCurrentWorkstationName = CopyNextWorkstationName(store, NULL);
if (gCurrentWorkstationName) {
gServiceRegistration = MyDNSServiceAlloc();
RegisterWorkstationService(gServiceRegistration, gCurrentWorkstationName);
}
}
static CFRunLoopSourceRef
InstallComputerNameMonitor(void)
{
CFRunLoopSourceRef runLoopSource = NULL;
SCDynamicStoreContext context = { 0, NULL, NULL, NULL, NULL };
SCDynamicStoreRef store = SCDynamicStoreCreate(NULL, CFSTR("DirectoryService"), ComputerNameChangedCallBack, &context);
if (store) {
CFStringRef computerNameKey = SCDynamicStoreKeyCreateComputerName(NULL);
CFMutableArrayRef keys = CFArrayCreateMutable(NULL, 1, &kCFTypeArrayCallBacks);
CFArrayAppendValue(keys, computerNameKey);
if (SCDynamicStoreSetNotificationKeys(store, keys, NULL)) {
runLoopSource = SCDynamicStoreCreateRunLoopSource(NULL, store, 0);
if (runLoopSource) {
CFRunLoopAddSource(CFRunLoopGetCurrent(), runLoopSource, kCFRunLoopCommonModes);
}
}
CFRelease(computerNameKey);
CFRelease(keys);
CFRelease(store);
}
return runLoopSource;
}
static void DisposeRegistration(void);
void DisposeRegistration(void)
{
if (gNameMonitorSource) {
CFRunLoopSourceInvalidate(gNameMonitorSource);
CFRelease(gNameMonitorSource);
gNameMonitorSource = NULL;
}
if (gServiceRegistration) {
MyDNSServiceFree(gServiceRegistration);
gServiceRegistration = NULL;
}
if (gCurrentWorkstationName) {
CFRelease(gCurrentWorkstationName);
gCurrentWorkstationName = NULL;
}
}
int32_t
WorkstationServiceRegister(void)
{
DNSServiceErrorType error = kDNSServiceErr_NoError;
CFStringRef workstationName = NULL;
workstationName = CopyNextWorkstationName(NULL, NULL);
if ( workstationName == NULL )
return kDNSServiceErr_NoSuchName;
if ( gCurrentWorkstationName != NULL )
{
if ( CFStringCompare(workstationName, gCurrentWorkstationName, 0) == kCFCompareEqualTo )
{
CFRelease( workstationName );
return kDNSServiceErr_NoError;
}
DisposeRegistration();
}
gCurrentWorkstationName = workstationName;
gNameMonitorSource = InstallComputerNameMonitor();
gServiceRegistration = MyDNSServiceAlloc();
error = RegisterWorkstationService(gServiceRegistration, gCurrentWorkstationName);
if (error != kDNSServiceErr_NoError)
{
DbgLog( kLogHandler, "WorkstationServiceRegister(): received error from RegisterWorkstationService: %l", error );
DisposeRegistration();
}
return error;
}
int32_t
WorkstationServiceUnregister(void)
{
DisposeRegistration();
return 0;
}