diff --git a/code/DMSShortcodeHandler.php b/code/DMSShortcodeHandler.php
index 452814f..e55c454 100644
--- a/code/DMSShortcodeHandler.php
+++ b/code/DMSShortcodeHandler.php
@@ -5,28 +5,35 @@
  */
 class DMSShortcodeHandler {
 
-	public static function handle($arguments, $content = null, $parser = null) {
-		$linkText = null;
+	public static function handle(
+		$arguments, $content, ShortcodeParser $parser, $tag, array $extra = array()
+	) {
+		if(!empty($arguments['id'])) {
+			$document = DMSDocument::get()->byID($arguments['id']);
 
-		if (!empty($arguments['id'])) {
-			$document = DMSDocument::get()->filter(array('ID' => $arguments['id']))->First();
-			if ($document && !$document->isHidden()) {
-				if (!empty($content)) {
-					$linkText = sprintf('<a href="%s">%s</a>', $document->Link(), $parser->parse($content));
+			if($document && !$document->isHidden()) {
+				if($content) {
+					return sprintf(
+						'<a href="%s">%s</a>', $document->Link(), $parser->parse($content)
+					);
 				} else {
-					$linkText = $document->Link();
+					if(isset($extra['element'])) {
+						$extra['element']->setAttribute('data-ext', $document->getExtension());
+						$extra['element']->setAttribute('data-size', $document->getFileSizeFormatted());
+					}
+
+					return $document->Link();
 				}
 			}
 		}
 
-		if (empty($linkText)) {
-			$errorPage = ErrorPage::get()->filter(array('ErrorCode' => '404'))->First();
-			if ($errorPage) {
-				$linkText = $errorPage->Link();
-			}
+		$error = ErrorPage::get()->filter('ErrorCode', '404')->First();
+
+		if($error) {
+			return $error->Link();
 		}
 
-		return $linkText;
+		return '';
 	}
 
 }
diff --git a/tests/DMSShortcodeTest.php b/tests/DMSShortcodeTest.php
new file mode 100644
index 0000000..4cbe864
--- /dev/null
+++ b/tests/DMSShortcodeTest.php
@@ -0,0 +1,23 @@
+<?php
+/**
+ * Tests DMS shortcode linking functionality.
+ */
+class DMSShortcodeTest extends SapphireTest {
+
+	public function testShortcodeOperation() {
+		$file = 'dms/tests/DMS-test-lorum-file.pdf';
+		$document = DMS::inst()->storeDocument($file);
+
+		$result = ShortcodeParser::get('default')->parse(sprintf(
+			'<p><a href="[dms_document_link,id=\'%d\']">Document</a></p>', $document->ID
+		));
+
+		$value = Injector::inst()->create('HTMLValue', $result);
+		$link = $value->query('//a')->item(0);
+
+		$this->assertEquals("/dmsdocument/$document->ID", $link->getAttribute('href'));
+		$this->assertEquals($document->getExtension(), $link->getAttribute('data-ext'));
+		$this->assertEquals($document->getFileSizeFormatted(), $link->getAttribute('data-size'));
+	}
+
+}